diff --git a/crates/forge_ci/src/jobs/bounty_job.rs b/crates/forge_ci/src/jobs/bounty_job.rs new file mode 100644 index 0000000000000000000000000000000000000000..bb0becca26ef61b993bdd9c58c91b76749e1537c --- /dev/null +++ b/crates/forge_ci/src/jobs/bounty_job.rs @@ -0,0 +1,74 @@ +//! Jobs for the bounty management workflow (v2). +//! +//! v2 uses a state-reconciliation model: each job fetches the full current +//! state of an issue or PR from GitHub, computes the desired label set from +//! the rules engine, diffs current vs desired, and applies the minimal patch. +//! +//! Three entry points: +//! - `sync-all-issues.ts` — fetches all open issues with any bounty label and +//! reconciles their label sets in one pass. Runs on a schedule and on label +//! events. +//! - `sync-pr.ts --pr N` — propagates labels from linked issues to the PR; +//! handles the rewarded lifecycle on merge. + +use gh_workflow::*; + +const SCRIPTS_DIR: &str = ".github/scripts/bounty/src"; +const TSX: &str = "npx tsx"; + +/// Returns a checkout step — required before script invocation. +fn checkout_step() -> Step { + Step::new("Checkout").uses("actions", "checkout", "v6") +} + +/// Builds a three-step job: checkout + npm install + a single script +/// invocation. +fn sync_job(job_name: &str, script: &str, args: String) -> Job { + let cmd = format!("{TSX} {SCRIPTS_DIR}/{script} {args}"); + Job::new(job_name) + .add_step(checkout_step()) + .add_step(Step::new("Install npm packages").run("npm install")) + .add_step(Step::new("Sync bounty labels").run(cmd)) +} + +/// Creates a job that syncs bounty labels across all open issues that carry +/// any bounty label. +/// +/// Fetches every open issue with a "bounty" label prefix, computes the desired +/// state for each, and applies the minimal patch in a single pass. +/// +/// Triggered on: issues labeled/unlabeled/assigned/unassigned, and on schedule. +pub fn sync_all_issues_job() -> Job { + let cmd = format!( + "{TSX} {SCRIPTS_DIR}/sync-all-issues.ts \ + --repo ${{{{ github.repository }}}} \ + --token ${{{{ secrets.GITHUB_TOKEN }}}} \ + --execute" + ); + Job::new("Sync all bounty issues") + .add_step(checkout_step()) + .add_step(Step::new("Install npm packages").run("npm install")) + .add_step(Step::new("Sync all bounty labels").run(cmd)) + .permissions(Permissions::default().issues(Level::Write)) +} + +/// Creates a job that propagates bounty labels from linked issues to the PR +/// and handles the rewarded lifecycle when the PR is merged. +/// +/// Triggered on: pull_request opened/edited/reopened, pull_request_target +/// closed. +pub fn sync_pr_job() -> Job { + sync_job( + "Sync PR bounty labels", + "sync-pr.ts", + "--pr ${{ github.event.pull_request.number }} \ + --repo ${{ github.repository }} \ + --token ${{ secrets.GITHUB_TOKEN }}" + .to_string(), + ) + .permissions( + Permissions::default() + .issues(Level::Write) + .pull_requests(Level::Write), + ) +} diff --git a/crates/forge_ci/src/jobs/build.rs b/crates/forge_ci/src/jobs/build.rs new file mode 100644 index 0000000000000000000000000000000000000000..723f65d24482851ad63850ff11a2240de059efeb --- /dev/null +++ b/crates/forge_ci/src/jobs/build.rs @@ -0,0 +1,9 @@ +use gh_workflow::*; + +use crate::jobs::release_build_job::release_build_job; + +/// Create a build job for drafts +pub fn create_build_release_job_for_publishing() -> Job { + release_build_job("${{ github.event.release.tag_name }}", true) +} + diff --git a/crates/forge_ci/src/jobs/draft_release_update_job.rs b/crates/forge_ci/src/jobs/draft_release_update_job.rs new file mode 100644 index 0000000000000000000000000000000000000000..ff5eb1513c79502ae7e7c417549e096e347588d0 --- /dev/null +++ b/crates/forge_ci/src/jobs/draft_release_update_job.rs @@ -0,0 +1,21 @@ +use gh_workflow::*; + +/// Create a job to update the release draft +pub fn draft_release_update_job() -> Job { + Job::new("update_release_draft") + .add_step( + Step::new("Auto Labeler") + .uses("release-drafter", "release-drafter/autolabeler", "v7") + .if_condition(Expression::new( + "github.event_name == 'pull_request_target'", + )) + .env(("GITHUB_TOKEN", "${{ secrets.GITHUB_TOKEN }}")) + .add_with(("config-name", "release-drafter.yml")), + ) + .add_step( + Step::new("Release Drafter") + .uses("release-drafter", "release-drafter", "v7") + .env(("GITHUB_TOKEN", "${{ secrets.GITHUB_TOKEN }}")) + .add_with(("config-name", "release-drafter.yml")), + ) +} diff --git a/crates/forge_ci/src/jobs/label_sync_job.rs b/crates/forge_ci/src/jobs/label_sync_job.rs new file mode 100644 index 0000000000000000000000000000000000000000..aeba168d11b44f3c211b93b81b4e30bc5af7f4a1 --- /dev/null +++ b/crates/forge_ci/src/jobs/label_sync_job.rs @@ -0,0 +1,20 @@ +use gh_workflow::*; + +/// Create a job to sync GitHub labels +pub fn label_sync_job() -> Job { + Job::new("label-sync") + .permissions( + Permissions::default() + .issues(Level::Write) + ) + .add_step( + Step::new("Checkout Code").uses("actions", "checkout", "v6") + .name("Checkout") + ) + .add_step( + Step::new("Sync Labels").run( + "npx github-label-sync \\\n --access-token ${{ secrets.GITHUB_TOKEN }} \\\n --labels \".github/labels.json\" \\\n ${{ github.repository }}" + ) + .name("Sync labels") + ) +} diff --git a/crates/forge_ci/src/jobs/lint.rs b/crates/forge_ci/src/jobs/lint.rs new file mode 100644 index 0000000000000000000000000000000000000000..b3f95b490046753a5deccc7597bafd93bd2dc69a --- /dev/null +++ b/crates/forge_ci/src/jobs/lint.rs @@ -0,0 +1,68 @@ +//! Shared lint commands for CI workflows + +/// Build a cargo command from parts +fn cargo_cmd(parts: &[&str]) -> String { + parts.join(" ") +} + +/// Base parts for fmt commands +fn fmt_base() -> Vec<&'static str> { + vec!["cargo", "+nightly", "fmt", "--all"] +} + +/// Base parts for clippy commands (shared across all clippy invocations). +fn clippy_base() -> Vec<&'static str> { + vec![ + "cargo", + "+nightly", + "clippy", + "--all-features", + "--workspace", + ] +} + +/// Build a cargo fmt command +pub fn fmt_cmd(fix: bool) -> String { + let mut parts = fmt_base(); + if !fix { + parts.push("--check"); + } + cargo_cmd(&parts) +} + +/// Build a cargo clippy command that checks all targets for general warnings. +pub fn clippy_cmd(fix: bool) -> String { + let mut parts = clippy_base(); + parts.push("--all-targets"); + + if fix { + parts.extend(["--fix", "--allow-dirty"]); + } + + parts.extend(["--", "-D", "warnings"]); + + cargo_cmd(&parts) +} + +/// Build a cargo clippy command for UTF-8 and indexing safety lints. +/// +/// Excludes test code by omitting `--all-targets`. +pub fn clippy_string_safety_cmd(fix: bool) -> String { + let mut parts = clippy_base(); + + if fix { + parts.extend(["--fix", "--allow-dirty"]); + } + + parts.extend([ + "--", + "-D", + "clippy::string_slice", + "-D", + "clippy::indexing_slicing", + "-D", + "clippy::disallowed_methods", + ]); + + cargo_cmd(&parts) +} diff --git a/crates/forge_ci/src/jobs/mod.rs b/crates/forge_ci/src/jobs/mod.rs new file mode 100644 index 0000000000000000000000000000000000000000..09731ab99d67970ad80fe8e4e860baab57a3aafb --- /dev/null +++ b/crates/forge_ci/src/jobs/mod.rs @@ -0,0 +1,21 @@ +//! Jobs for CI workflows + +mod bounty_job; +mod draft_release_update_job; +mod label_sync_job; +mod lint; +mod release_build_job; +mod release_draft; +mod release_draft_pr; +mod release_homebrew; +mod release_npm; + +pub use bounty_job::*; +pub use draft_release_update_job::*; +pub use label_sync_job::*; +pub use lint::*; +pub use release_build_job::*; +pub use release_draft::*; +pub use release_draft_pr::*; +pub use release_homebrew::*; +pub use release_npm::*; diff --git a/crates/forge_ci/src/jobs/release_build_job.rs b/crates/forge_ci/src/jobs/release_build_job.rs new file mode 100644 index 0000000000000000000000000000000000000000..edd3a0662cd91e13bde46fd211e334e9373a69a4 --- /dev/null +++ b/crates/forge_ci/src/jobs/release_build_job.rs @@ -0,0 +1,103 @@ +use derive_setters::Setters; +use gh_workflow::*; + +use crate::release_matrix::ReleaseMatrix; +use crate::steps::setup_protoc; + +#[derive(Clone, Default, Setters)] +#[setters(strip_option, into)] +pub struct ReleaseBuilderJob { + // Required to burn into the binary + pub version: String, + + // When provide the generated release will be uploaded + pub release_id: Option, +} + +impl ReleaseBuilderJob { + pub fn new(version: impl AsRef) -> Self { + Self { version: version.as_ref().to_string(), release_id: None } + } + + pub fn into_job(self) -> Job { + self.into() + } +} + +impl From for Job { + fn from(value: ReleaseBuilderJob) -> Job { + let mut job = Job::new("build-release") + .strategy(Strategy { + fail_fast: None, + max_parallel: None, + matrix: Some(ReleaseMatrix::default().into()), + }) + .runs_on("${{ matrix.os }}") + .permissions( + Permissions::default() + .contents(Level::Write) + .pull_requests(Level::Write), + ) + .add_step(Step::new("Checkout Code").uses("actions", "checkout", "v6")) + // Install protobuf compiler for non-cross builds + // Cross builds install protoc via Cross.toml pre-build commands + .add_step( + setup_protoc().if_condition(Expression::new("${{ matrix.cross == 'false' }}")), + ) + // Install Rust with cross-compilation target + .add_step( + Step::new("Setup Cross Toolchain") + .uses("taiki-e", "setup-cross-toolchain-action", "v1") + .with(("target", "${{ matrix.target }}")) + .if_condition(Expression::new("${{ matrix.cross == 'false' }}")), + ) + // Explicitly add the target to ensure it's available + .add_step( + Step::new("Add Rust target") + .run("rustup target add ${{ matrix.target }}") + .if_condition(Expression::new("${{ matrix.cross == 'false' }}")), + ) + // Build add link flags + .add_step( + Step::new("Set Rust Flags") + .run(r#"echo "RUSTFLAGS=-C target-feature=+crt-static" >> $GITHUB_ENV"#) + .if_condition(Expression::new( + "!(contains(matrix.target, '-unknown-linux-') || contains(matrix.target, '-android'))", + )), + ) + // Build release binary + // Note: protoc is installed via: + // - arduino/setup-protoc action for non-cross builds + // - Cross.toml pre-build commands for cross builds (apt-get install protobuf-compiler) + .add_step( + Step::new("Build Binary") + .uses("ClementTsang", "cargo-action", "v0.0.7") + .add_with(("command", "build --release")) + .add_with(("args", "--target ${{ matrix.target }}")) + .add_with(("use-cross", "${{ matrix.cross }}")) + .add_with(("cross-version", "0.2.5")) + .add_env(("RUSTFLAGS", "${{ env.RUSTFLAGS }}")) + .add_env(("POSTHOG_API_SECRET", "${{secrets.POSTHOG_API_SECRET}}")) + .add_env(("APP_VERSION", value.version.to_string())), + ); + + if let Some(release_id) = value.release_id { + job = job + // Rename binary to target name + .add_step( + Step::new("Copy Binary") + .run("cp ${{ matrix.binary_path }} ${{ matrix.binary_name }}"), + ) + // Upload to the generated github release id + .add_step( + Step::new("Upload to Release") + .uses("xresloader", "upload-to-github-release", "v1") + .add_with(("release_id", release_id)) + .add_with(("file", "${{ matrix.binary_name }}")) + .add_with(("overwrite", "true")), + ); + } + + job + } +} diff --git a/crates/forge_ci/src/jobs/release_draft.rs b/crates/forge_ci/src/jobs/release_draft.rs new file mode 100644 index 0000000000000000000000000000000000000000..7565d85d8ed6247b74e0cb558eb17e36f038bfd5 --- /dev/null +++ b/crates/forge_ci/src/jobs/release_draft.rs @@ -0,0 +1,33 @@ +use gh_workflow::*; +use indexmap::indexmap; + +/// Create a draft release job for GitHub Actions +pub fn create_draft_release_job(build_job_id: &str) -> Job { + Job::new("draft_release") + .name("Draft Release") + .add_needs(build_job_id) + .cond(Expression::new( + "github.event_name == 'push' && github.ref == 'refs/heads/main'", + )) + // This job only runs on push to main, not needed for release events + .permissions( + Permissions::default() + .contents(Level::Write) + .pull_requests(Level::Write), + ) + .add_step(Step::new("Checkout Code").uses("actions", "checkout", "v6")) + .add_step( + Step::new("Draft Release").uses("release-drafter", "release-drafter", "v7") + .id("create_release") + .env(("GITHUB_TOKEN", "${{ secrets.GITHUB_TOKEN }}")) + .with(("config-name", "release-drafter.yml")), + ) + .add_step( + Step::new("Export Outputs").run("echo \"crate_release_id=${{ steps.create_release.outputs.id }}\" >> $GITHUB_OUTPUT && echo \"crate_release_name=${{ steps.create_release.outputs.tag_name }}\" >> $GITHUB_OUTPUT") + .id("set_output"), + ) + .outputs(indexmap! { + "crate_release_name".to_string() => "${{ steps.set_output.outputs.crate_release_name }}".to_string(), + "crate_release_id".to_string() => "${{ steps.set_output.outputs.crate_release_id }}".to_string() + }) +} diff --git a/crates/forge_ci/src/jobs/release_draft_pr.rs b/crates/forge_ci/src/jobs/release_draft_pr.rs new file mode 100644 index 0000000000000000000000000000000000000000..d9b5cca937543b70a08d19c0f92a072661e6bc71 --- /dev/null +++ b/crates/forge_ci/src/jobs/release_draft_pr.rs @@ -0,0 +1,22 @@ +use gh_workflow::*; +use indexmap::indexmap; + +/// Create a draft release job for GitHub Actions that runs on PRs +pub fn create_draft_release_pr_job() -> Job { + Job::new("draft_release_pr") + .name("Draft Release for PR") + .cond(Expression::new( + "github.event_name == 'pull_request' && contains(github.event.pull_request.labels.*.name, 'ci: build all targets')", + )) + .add_step(Step::new("Checkout Code").uses("actions", "checkout", "v6")) + .add_step( + Step::new("Set Release Version").run( + r#"echo "crate_release_name=pr-build-${{ github.event.number }}" >> $GITHUB_OUTPUT && echo "crate_release_id=pr-build-${{ github.event.number }}" >> $GITHUB_OUTPUT"#, + ) + .id("set_output"), + ) + .outputs(indexmap! { + "crate_release_name".to_string() => "${{ steps.set_output.outputs.crate_release_name }}".to_string(), + "crate_release_id".to_string() => "${{ steps.set_output.outputs.crate_release_id }}".to_string() + }) +} diff --git a/crates/forge_ci/src/jobs/release_homebrew.rs b/crates/forge_ci/src/jobs/release_homebrew.rs new file mode 100644 index 0000000000000000000000000000000000000000..12b081715256841dd725228b218325198b8fb5b4 --- /dev/null +++ b/crates/forge_ci/src/jobs/release_homebrew.rs @@ -0,0 +1,16 @@ +use gh_workflow::*; + +/// Create a homebrew release job +pub fn release_homebrew_job() -> Job { + Job::new("homebrew_release") + .add_step( + Step::new("Checkout Code").uses("actions", "checkout", "v6") + .add_with(("repository", "antinomyhq/homebrew-code-forge")) + .add_with(("ref", "main")) + .add_with(("token", "${{ secrets.HOMEBREW_ACCESS }}")), + ) + // Make script executable and run it with token + .add_step( + Step::new("Update Homebrew Formula").run("GITHUB_TOKEN=\"${{ secrets.HOMEBREW_ACCESS }}\" ./update-formula.sh ${{ github.event.release.tag_name }}"), + ) +} diff --git a/crates/forge_ci/src/jobs/release_npm.rs b/crates/forge_ci/src/jobs/release_npm.rs new file mode 100644 index 0000000000000000000000000000000000000000..7a789a319738880d123b1311c97d5a6f7f02bbdf --- /dev/null +++ b/crates/forge_ci/src/jobs/release_npm.rs @@ -0,0 +1,35 @@ +use gh_workflow::*; +use serde_json::Value; + +/// Create an NPM release job using matrix strategy for multiple repositories +pub fn release_npm_job() -> Job { + let matrix = create_npm_matrix(); + + Job::new("npm_release") + .strategy(Strategy { fail_fast: None, max_parallel: None, matrix: Some(matrix) }) + .add_step( + Step::new("Checkout Code") + .uses("actions", "checkout", "v6") + .add_with(("repository", "${{ matrix.repository }}")) + .add_with(("ref", "main")) + .add_with(("token", "${{ secrets.NPM_ACCESS }}")), + ) + // Make script executable and run it with token + .add_step( + Step::new("Update NPM Package") + .run("./update-package.sh ${{ github.event.release.tag_name }}") + .add_env(("AUTO_PUSH", "true")) + .add_env(("CI", "true")) + .add_env(("NPM_TOKEN", "${{ secrets.NPM_TOKEN }}")), + ) +} + +/// Creates a matrix Value for NPM repositories +fn create_npm_matrix() -> Value { + serde_json::json!({ + "repository": [ + "antinomyhq/npm-code-forge", + "antinomyhq/npm-forgecode" + ] + }) +} diff --git a/crates/forge_ci/src/lib.rs b/crates/forge_ci/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..22112e62f2e53bd01628a3f44e337394cbbb2be6 --- /dev/null +++ b/crates/forge_ci/src/lib.rs @@ -0,0 +1,4 @@ +mod jobs; +mod release_matrix; +pub mod steps; +pub mod workflows; diff --git a/crates/forge_ci/src/release_matrix.rs b/crates/forge_ci/src/release_matrix.rs new file mode 100644 index 0000000000000000000000000000000000000000..00b6a25ad909e5d609fd239cf57780512ee6c262 --- /dev/null +++ b/crates/forge_ci/src/release_matrix.rs @@ -0,0 +1,100 @@ +use serde::Serialize; +use serde_json::Value; + +/// Matrix entry for build targets +#[derive(Serialize, Clone)] +pub struct MatrixEntry { + pub os: &'static str, + pub target: &'static str, + pub binary_name: &'static str, + pub binary_path: &'static str, + pub cross: &'static str, +} + +#[derive(Clone)] +pub struct ReleaseMatrix(Vec); + +impl Default for ReleaseMatrix { + /// Returns a vector of all build matrix entries + fn default() -> Self { + ReleaseMatrix(vec![ + MatrixEntry { + os: "ubuntu-latest", + target: "x86_64-unknown-linux-musl", + binary_name: "forge-x86_64-unknown-linux-musl", + binary_path: "target/x86_64-unknown-linux-musl/release/forge", + cross: "true", + }, + MatrixEntry { + os: "ubuntu-latest", + target: "aarch64-unknown-linux-musl", + binary_name: "forge-aarch64-unknown-linux-musl", + binary_path: "target/aarch64-unknown-linux-musl/release/forge", + cross: "true", + }, + MatrixEntry { + os: "ubuntu-latest", + target: "x86_64-unknown-linux-gnu", + binary_name: "forge-x86_64-unknown-linux-gnu", + binary_path: "target/x86_64-unknown-linux-gnu/release/forge", + cross: "false", + }, + MatrixEntry { + os: "ubuntu-latest", + target: "aarch64-unknown-linux-gnu", + binary_name: "forge-aarch64-unknown-linux-gnu", + binary_path: "target/aarch64-unknown-linux-gnu/release/forge", + cross: "true", + }, + MatrixEntry { + os: "macos-latest", + target: "x86_64-apple-darwin", + binary_name: "forge-x86_64-apple-darwin", + binary_path: "target/x86_64-apple-darwin/release/forge", + cross: "false", + }, + MatrixEntry { + os: "macos-latest", + target: "aarch64-apple-darwin", + binary_name: "forge-aarch64-apple-darwin", + binary_path: "target/aarch64-apple-darwin/release/forge", + cross: "false", + }, + MatrixEntry { + os: "windows-latest", + target: "x86_64-pc-windows-msvc", + binary_name: "forge-x86_64-pc-windows-msvc.exe", + binary_path: "target/x86_64-pc-windows-msvc/release/forge.exe", + cross: "false", + }, + MatrixEntry { + os: "windows-latest", + target: "aarch64-pc-windows-msvc", + binary_name: "forge-aarch64-pc-windows-msvc.exe", + binary_path: "target/aarch64-pc-windows-msvc/release/forge.exe", + cross: "false", + }, + MatrixEntry { + os: "ubuntu-latest", + target: "aarch64-linux-android", + binary_name: "forge-aarch64-linux-android", + binary_path: "target/aarch64-linux-android/release/forge", + cross: "true", + }, + ]) + } +} + +impl ReleaseMatrix { + pub fn entries(&self) -> Vec { + self.0.clone() + } +} + +impl From for Value { + fn from(value: ReleaseMatrix) -> Self { + serde_json::json!({ + "include": value.entries() + }) + } +} diff --git a/crates/forge_ci/src/steps/mod.rs b/crates/forge_ci/src/steps/mod.rs new file mode 100644 index 0000000000000000000000000000000000000000..123ec5d55b4d2f018b28606436661c419f214998 --- /dev/null +++ b/crates/forge_ci/src/steps/mod.rs @@ -0,0 +1,3 @@ +mod setup_protoc; + +pub use setup_protoc::setup_protoc; diff --git a/crates/forge_ci/src/steps/setup_protoc.rs b/crates/forge_ci/src/steps/setup_protoc.rs new file mode 100644 index 0000000000000000000000000000000000000000..381d4912a19d2f595e04b3fd22c30fdb7c433bdd --- /dev/null +++ b/crates/forge_ci/src/steps/setup_protoc.rs @@ -0,0 +1,11 @@ +use gh_workflow::*; + +/// Creates a step to setup the Protobuf compiler. +/// +/// This step is reusable across all CI workflows that need protobuf +/// compilation. +pub fn setup_protoc() -> Step { + Step::new("Setup Protobuf Compiler") + .uses("arduino", "setup-protoc", "v3") + .with(("repo-token", "${{ secrets.GITHUB_TOKEN }}")) +} diff --git a/crates/forge_ci/src/workflows/autofix.rs b/crates/forge_ci/src/workflows/autofix.rs new file mode 100644 index 0000000000000000000000000000000000000000..102ae5019bb861856e7c88f9d089f493734cb12d --- /dev/null +++ b/crates/forge_ci/src/workflows/autofix.rs @@ -0,0 +1,57 @@ +use gh_workflow::generate::Generate; +use gh_workflow::toolchain::Component; +use gh_workflow::*; + +use crate::jobs; +use crate::steps::setup_protoc; + +/// Generate the autofix workflow +pub fn generate_autofix_workflow() { + let lint_fix_job = Job::new("Lint Fix") + .permissions(Permissions::default().contents(Level::Read)) + .add_step(Step::new("Checkout Code").uses("actions", "checkout", "v6")) + .add_step(Step::new("Install SQLite").run("sudo apt-get install -y libsqlite3-dev")) + .add_step(setup_protoc()) + .add_step( + Step::toolchain() + .add_nightly() + .add_component(Component::Clippy) + .add_component(Component::Rustfmt), + ) + .add_step(Step::new("Cargo Fmt").run(jobs::fmt_cmd(true))) + .add_step(Step::new("Cargo Clippy").run(jobs::clippy_cmd(true))) + .add_step( + Step::new("Cargo Clippy String Safety").run(jobs::clippy_string_safety_cmd(false)), + ) + .add_step(Step::new("Autofix").uses( + "autofix-ci", + "action", + "7a166d7532b277f34e16238930461bf77f9d7ed8", + )); + + let events = Event::default() + .push(Push::default().add_branch("main")) + .pull_request( + PullRequest::default() + .add_type(PullRequestType::Opened) + .add_type(PullRequestType::Synchronize) + .add_type(PullRequestType::Reopened) + .add_branch("main"), + ); + + let workflow = Workflow::default() + .name("autofix.ci") + .add_env(RustFlags::deny("warnings")) + .on(events) + .concurrency( + Concurrency::default() + .group("autofix-${{github.ref}}") + .cancel_in_progress(false), + ) + .add_job("lint", lint_fix_job); + + Generate::new(workflow) + .name("autofix.yml") + .generate() + .unwrap(); +} diff --git a/crates/forge_ci/src/workflows/bounty.rs b/crates/forge_ci/src/workflows/bounty.rs new file mode 100644 index 0000000000000000000000000000000000000000..99854fc4ffdd69e25db9e6d7e6c16b86ec3eb226 --- /dev/null +++ b/crates/forge_ci/src/workflows/bounty.rs @@ -0,0 +1,47 @@ +use gh_workflow::generate::Generate; +use gh_workflow::*; + +use crate::jobs::{sync_all_issues_job, sync_pr_job}; + +/// Generate the bounty management workflow (v2). +/// +/// Two jobs cover the full bounty lifecycle: +/// - `sync-all-issues`: fetches all open issues with any bounty label and +/// reconciles their label sets in one pass. Triggered on label/assignment +/// events and daily on a schedule. +/// - `sync-pr`: propagates bounty value labels from linked issues to the PR on +/// open/edit, and applies the rewarded lifecycle on merge. +pub fn generate_bounty_workflow() { + let events = Event::default() + .pull_request( + PullRequest::default() + .add_type(PullRequestType::Opened) + .add_type(PullRequestType::Edited) + .add_type(PullRequestType::Reopened), + ) + .pull_request_target(PullRequestTarget::default().add_type(PullRequestType::Closed)) + .issues( + Issues::default() + .add_type(IssuesType::Assigned) + .add_type(IssuesType::Unassigned) + .add_type(IssuesType::Labeled) + .add_type(IssuesType::Unlabeled), + ) + .add_cron_schedule("0 2 * * *"); + + let workflow = Workflow::default() + .name("Bounty Management") + .on(events) + .permissions( + Permissions::default() + .issues(Level::Write) + .pull_requests(Level::Write), + ) + .add_job("sync-all-issues", sync_all_issues_job()) + .add_job("sync-pr", sync_pr_job()); + + Generate::new(workflow) + .name("bounty.yml") + .generate() + .unwrap(); +} diff --git a/crates/forge_ci/src/workflows/ci.rs b/crates/forge_ci/src/workflows/ci.rs new file mode 100644 index 0000000000000000000000000000000000000000..2b5d81174fa35dbbdc12f5276b60e864b343877b --- /dev/null +++ b/crates/forge_ci/src/workflows/ci.rs @@ -0,0 +1,82 @@ +use gh_workflow::generate::Generate; +use gh_workflow::*; + +use crate::jobs::{self, ReleaseBuilderJob}; +use crate::steps::setup_protoc; + +/// Generate the main CI workflow +pub fn generate_ci_workflow() { + // Create a basic build job for CI with coverage + let build_job = Job::new("Build and Test") + .permissions(Permissions::default().contents(Level::Read)) + .add_step(Step::new("Checkout Code").uses("actions", "checkout", "v6")) + .add_step(setup_protoc()) + .add_step(Step::toolchain().add_stable()) + .add_step(Step::new("Install cargo-llvm-cov").run("cargo install cargo-llvm-cov")) + .add_step( + Step::new("Generate coverage") + .run("cargo llvm-cov --all-features --workspace --lcov --output-path lcov.info"), + ); + + // Create a performance test job to ensure zsh rprompt stays fast + let perf_test_job = Job::new("zsh-rprompt-performance") + .name("Performance: zsh rprompt") + .permissions(Permissions::default().contents(Level::Read)) + .add_step(Step::new("Checkout Code").uses("actions", "checkout", "v6")) + .add_step(setup_protoc()) + .add_step(Step::toolchain().add_stable()) + .add_step( + Step::new("Run performance benchmark") + .run("./scripts/benchmark.sh --threshold 60 zsh rprompt"), + ); + + let draft_release_job = jobs::create_draft_release_job("build"); + let draft_release_pr_job = jobs::create_draft_release_pr_job(); + let events = Event::default() + .push(Push::default().add_branch("main").add_tag("v*")) + .pull_request( + PullRequest::default() + .add_type(PullRequestType::Opened) + .add_type(PullRequestType::Synchronize) + .add_type(PullRequestType::Reopened) + .add_type(PullRequestType::Labeled) + .add_branch("main"), + ); + let build_release_pr_job = + ReleaseBuilderJob::new("${{ needs.draft_release_pr.outputs.crate_release_name }}") + .into_job() + .add_needs("draft_release_pr") + .cond(Expression::new( + [ + "github.event_name == 'pull_request'", + "contains(github.event.pull_request.labels.*.name, 'ci: build all targets')", + ] + .join(" && "), + )); + let build_release_job = + ReleaseBuilderJob::new("${{ needs.draft_release.outputs.crate_release_name }}") + .release_id("${{ needs.draft_release.outputs.crate_release_id }}") + .into_job() + .add_needs("draft_release") + .cond(Expression::new( + [ + "github.event_name == 'push'", + "github.ref == 'refs/heads/main'", + ] + .join(" && "), + )); + let workflow = Workflow::default() + .name("ci") + .add_env(RustFlags::deny("warnings")) + .on(events) + .concurrency(Concurrency::default().group("${{ github.workflow }}-${{ github.ref }}")) + .add_env(("OPENROUTER_API_KEY", "${{secrets.OPENROUTER_API_KEY}}")) + .add_job("build", build_job) + .add_job("zsh_rprompt_perf", perf_test_job) + .add_job("draft_release", draft_release_job) + .add_job("draft_release_pr", draft_release_pr_job) + .add_job("build_release", build_release_job) + .add_job("build_release_pr", build_release_pr_job); + + Generate::new(workflow).name("ci.yml").generate().unwrap(); +} diff --git a/crates/forge_ci/src/workflows/labels.rs b/crates/forge_ci/src/workflows/labels.rs new file mode 100644 index 0000000000000000000000000000000000000000..7a72a6a93e7368961ac6fd85888680be68371edc --- /dev/null +++ b/crates/forge_ci/src/workflows/labels.rs @@ -0,0 +1,21 @@ +use gh_workflow::generate::Generate; +use gh_workflow::*; + +use crate::jobs::label_sync_job; + +/// Generate labels workflow +pub fn generate_labels_workflow() { + let labels_workflow = Workflow::default() + .name("Github Label Sync") + .on(Event { + push: Some(Push { branches: vec!["main".to_string()], ..Push::default() }), + ..Event::default() + }) + .permissions(Permissions::default().contents(Level::Write)) + .add_job("label-sync", label_sync_job()); + + Generate::new(labels_workflow) + .name("labels.yml") + .generate() + .unwrap(); +} diff --git a/crates/forge_ci/src/workflows/mod.rs b/crates/forge_ci/src/workflows/mod.rs new file mode 100644 index 0000000000000000000000000000000000000000..ff76efb7841a231f9a76982a61109d7ead125199 --- /dev/null +++ b/crates/forge_ci/src/workflows/mod.rs @@ -0,0 +1,17 @@ +//! Workflow definitions for CI/CD + +mod autofix; +mod bounty; +mod ci; +mod labels; +mod release_drafter; +mod release_publish; +mod stale; + +pub use autofix::*; +pub use bounty::*; +pub use ci::*; +pub use labels::*; +pub use release_drafter::*; +pub use release_publish::*; +pub use stale::*; diff --git a/crates/forge_ci/src/workflows/release_drafter.rs b/crates/forge_ci/src/workflows/release_drafter.rs new file mode 100644 index 0000000000000000000000000000000000000000..982d9a77e17820eab920d36024295fd34a58c855 --- /dev/null +++ b/crates/forge_ci/src/workflows/release_drafter.rs @@ -0,0 +1,36 @@ +use gh_workflow::generate::Generate; +use gh_workflow::*; + +use crate::jobs::draft_release_update_job; + +/// Generate release drafter workflow +pub fn generate_release_drafter_workflow() { + let release_drafter = Workflow::default() + .name("Release Drafter") + .on(Event { + push: Some(Push { branches: vec!["main".to_string()], ..Push::default() }), + pull_request_target: Some(PullRequestTarget { + types: vec![ + PullRequestType::Opened, + PullRequestType::Reopened, + PullRequestType::Synchronize, + PullRequestType::Labeled, + PullRequestType::Unlabeled, + PullRequestType::Closed, + ], + branches: vec!["main".to_string()], + }), + ..Event::default() + }) + .permissions( + Permissions::default() + .contents(Level::Write) + .pull_requests(Level::Write), + ) + .add_job("update_release_draft", draft_release_update_job()); + + Generate::new(release_drafter) + .name("release-drafter.yml") + .generate() + .unwrap(); +} diff --git a/crates/forge_ci/src/workflows/release_publish.rs b/crates/forge_ci/src/workflows/release_publish.rs new file mode 100644 index 0000000000000000000000000000000000000000..fbbd9fd834c91e998a56c0e1c0b672865a570c80 --- /dev/null +++ b/crates/forge_ci/src/workflows/release_publish.rs @@ -0,0 +1,32 @@ +use gh_workflow::generate::Generate; +use gh_workflow::*; + +use crate::jobs::{ReleaseBuilderJob, release_homebrew_job, release_npm_job}; + +/// Generate npm release workflow +pub fn release_publish() { + let release_build_job = ReleaseBuilderJob::new("${{ github.event.release.tag_name }}") + .release_id("${{ github.event.release.id }}"); + let npm_release_job = release_npm_job().add_needs("build_release"); + let homebrew_release_job = release_homebrew_job().add_needs("build_release"); + + let npm_workflow = Workflow::default() + .name("Multi Channel Release") + .on(Event { + release: Some(Release { types: vec![ReleaseType::Published] }), + ..Event::default() + }) + .permissions( + Permissions::default() + .contents(Level::Write) + .pull_requests(Level::Write), + ) + .add_job("build_release", release_build_job.into_job()) + .add_job("npm_release", npm_release_job) + .add_job("homebrew_release", homebrew_release_job); + + Generate::new(npm_workflow) + .name("release.yml") + .generate() + .unwrap(); +} diff --git a/crates/forge_ci/src/workflows/stale.rs b/crates/forge_ci/src/workflows/stale.rs new file mode 100644 index 0000000000000000000000000000000000000000..374d61bfee4212f33baa2c87df042ce6d7ac9306 --- /dev/null +++ b/crates/forge_ci/src/workflows/stale.rs @@ -0,0 +1,48 @@ +use gh_workflow::generate::Generate; +use gh_workflow::*; +use indexmap::indexmap; +use serde_json::json; + +/// Generate the stale issues and PRs workflow +pub fn generate_stale_workflow() { + let workflow = Workflow::default() + .name("Close Stale Issues and PR") + .on(Event::default().add_schedule(Schedule::new("0 * * * *"))) + .permissions( + Permissions::default() + .issues(Level::Write) + .pull_requests(Level::Write), + ) + .env(Env::from(indexmap! { + "DAYS_BEFORE_ISSUE_STALE".to_string() => json!("30"), + "DAYS_BEFORE_ISSUE_CLOSE".to_string() => json!("7"), + "DAYS_BEFORE_PR_STALE".to_string() => json!("5"), + "DAYS_BEFORE_PR_CLOSE".to_string() => json!("10"), + })) + .add_job( + "stale", + Job::new("Stale Issues") + .add_step( + Step::new("Mark Stale Issues").uses("actions", "stale", "v10") + .with(Input::from(indexmap! { + "stale-issue-label".to_string() => json!("state: inactive"), + "stale-pr-label".to_string() => json!("state: inactive"), + "stale-issue-message".to_string() => json!(r#"**Action required:** Issue inactive for ${{ env.DAYS_BEFORE_ISSUE_STALE }} days. +Status update or closure in ${{ env.DAYS_BEFORE_ISSUE_CLOSE }} days."#), + "close-issue-message".to_string() => json!("Issue closed after ${{ env.DAYS_BEFORE_ISSUE_CLOSE }} days of inactivity."), + "stale-pr-message".to_string() => json!(r#"**Action required:** PR inactive for ${{ env.DAYS_BEFORE_PR_STALE }} days. +Status update or closure in ${{ env.DAYS_BEFORE_PR_CLOSE }} days."#), + "close-pr-message".to_string() => json!("PR closed after ${{ env.DAYS_BEFORE_PR_CLOSE }} days of inactivity."), + "days-before-issue-stale".to_string() => json!("${{ env.DAYS_BEFORE_ISSUE_STALE }}"), + "days-before-issue-close".to_string() => json!("${{ env.DAYS_BEFORE_ISSUE_CLOSE }}"), + "days-before-pr-stale".to_string() => json!("${{ env.DAYS_BEFORE_PR_STALE }}"), + "days-before-pr-close".to_string() => json!("${{ env.DAYS_BEFORE_PR_CLOSE }}"), + })), + ), + ); + + Generate::new(workflow) + .name("stale.yml") + .generate() + .unwrap(); +} diff --git a/crates/forge_ci/tests/ci.rs b/crates/forge_ci/tests/ci.rs new file mode 100644 index 0000000000000000000000000000000000000000..4856d76093b0424ae5008de186a05ddb7e5b26ab --- /dev/null +++ b/crates/forge_ci/tests/ci.rs @@ -0,0 +1,36 @@ +use forge_ci::workflows as workflow; + +#[test] +fn generate() { + workflow::generate_ci_workflow(); +} + +#[test] +fn test_release_drafter() { + workflow::generate_release_drafter_workflow(); +} + +#[test] +fn test_release_workflow() { + workflow::release_publish(); +} + +#[test] +fn test_labels_workflow() { + workflow::generate_labels_workflow(); +} + +#[test] +fn test_stale_workflow() { + workflow::generate_stale_workflow(); +} + +#[test] +fn test_autofix_workflow() { + workflow::generate_autofix_workflow(); +} + +#[test] +fn test_bounty_workflow() { + workflow::generate_bounty_workflow(); +} diff --git a/crates/forge_config/src/auto_dump.rs b/crates/forge_config/src/auto_dump.rs new file mode 100644 index 0000000000000000000000000000000000000000..ebf4e2100174a3abb7277f7b628b2ec407e2a3f4 --- /dev/null +++ b/crates/forge_config/src/auto_dump.rs @@ -0,0 +1,25 @@ +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +/// The output format used when auto-dumping a conversation on task completion. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, fake::Dummy)] +#[serde(rename_all = "snake_case")] +pub enum AutoDumpFormat { + /// Dump as a JSON file + Json, + /// Dump as an HTML file + Html, +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::*; + + #[test] + fn test_auto_dump_format_variants() { + assert_eq!(AutoDumpFormat::Json, AutoDumpFormat::Json); + assert_eq!(AutoDumpFormat::Html, AutoDumpFormat::Html); + } +} diff --git a/crates/forge_config/src/compact.rs b/crates/forge_config/src/compact.rs new file mode 100644 index 0000000000000000000000000000000000000000..06240052eb020b48e3bf70c0693c699425b13b04 --- /dev/null +++ b/crates/forge_config/src/compact.rs @@ -0,0 +1,266 @@ +use std::time::Duration; + +use derive_setters::Setters; +use fake::Dummy; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +use crate::Percentage; + +/// Frequency at which forge checks for updates +#[derive(Default, Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq, fake::Dummy)] +#[serde(rename_all = "snake_case")] +pub enum UpdateFrequency { + Daily, + Weekly, + Never, + #[default] + Always, +} + +impl From for Duration { + fn from(val: UpdateFrequency) -> Self { + match val { + UpdateFrequency::Daily => Duration::from_secs(60 * 60 * 24), + UpdateFrequency::Weekly => Duration::from_secs(60 * 60 * 24 * 7), + UpdateFrequency::Never => Duration::MAX, + UpdateFrequency::Always => Duration::ZERO, + } + } +} + +/// Configuration for automatic forge updates +#[derive( + Debug, Clone, Serialize, Deserialize, Default, JsonSchema, Setters, PartialEq, fake::Dummy, +)] +#[setters(strip_option, into)] +pub struct Update { + /// How frequently forge checks for updates: daily, weekly, always, or never + pub frequency: Option, + /// Whether to automatically install updates without prompting + pub auto_update: Option, +} + +/// Configuration for automatic context compaction for all agents +#[derive(Debug, Clone, Serialize, Deserialize, Setters, JsonSchema, PartialEq)] +#[setters(strip_option, into)] +pub struct Compact { + /// Number of most recent messages to preserve during compaction. + /// These messages won't be considered for summarization. Works alongside + /// eviction_window - the more conservative limit (fewer messages to + /// compact) takes precedence. + #[serde(default)] + pub retention_window: usize, + + /// Maximum percentage of the context that can be summarized during + /// compaction. Valid values are between 0.0 and 1.0, where 0.0 means no + /// compaction and 1.0 allows summarizing all messages. Works alongside + /// retention_window - the more conservative limit (fewer messages to + /// compact) takes precedence. + #[serde(default)] + pub eviction_window: Percentage, + + /// Maximum number of tokens to keep after compaction + #[serde(skip_serializing_if = "Option::is_none")] + pub max_tokens: Option, + + /// Maximum number of tokens before triggering compaction. This acts as an + /// absolute cap and is combined with + /// `token_threshold_percentage` by taking the lower value. + #[serde(skip_serializing_if = "Option::is_none")] + pub token_threshold: Option, + + /// Maximum percentage of the model context window used to derive the token + /// threshold before triggering compaction. This is combined with + /// `token_threshold` by taking the lower value. + #[serde(skip_serializing_if = "Option::is_none")] + pub token_threshold_percentage: Option, + + /// Maximum number of conversation turns before triggering compaction + #[serde(skip_serializing_if = "Option::is_none")] + pub turn_threshold: Option, + + /// Maximum number of messages before triggering compaction + #[serde(skip_serializing_if = "Option::is_none")] + pub message_threshold: Option, + + /// Model ID to use for compaction, useful when compacting with a + /// cheaper/faster model. If not specified, the root level model will be + /// used. + #[serde(skip_serializing_if = "Option::is_none")] + pub model: Option, + + /// Whether to trigger compaction when the last message is from a user + #[serde(default, skip_serializing_if = "Option::is_none")] + pub on_turn_end: Option, +} + +impl Default for Compact { + fn default() -> Self { + Self::new() + } +} + +impl Compact { + /// Creates a new compaction configuration with all optional fields unset + pub fn new() -> Self { + Self { + max_tokens: None, + token_threshold: None, + token_threshold_percentage: None, + turn_threshold: None, + message_threshold: None, + model: None, + eviction_window: Percentage::new(0.2).unwrap(), + retention_window: 0, + on_turn_end: None, + } + } +} + +impl Dummy for Compact { + fn dummy_with_rng(_: &fake::Faker, rng: &mut R) -> Self { + use fake::Fake; + Self { + retention_window: fake::Faker.fake_with_rng(rng), + eviction_window: Percentage::from((0.0f64..=1.0f64).fake_with_rng::(rng)), + max_tokens: fake::Faker.fake_with_rng(rng), + token_threshold: fake::Faker.fake_with_rng(rng), + token_threshold_percentage: fake::Faker.fake_with_rng(rng), + turn_threshold: fake::Faker.fake_with_rng(rng), + message_threshold: fake::Faker.fake_with_rng(rng), + model: fake::Faker.fake_with_rng(rng), + on_turn_end: fake::Faker.fake_with_rng(rng), + } + } +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::*; + use crate::ForgeConfig; + use crate::reader::ConfigReader; + + #[test] + fn test_f64_eviction_window_round_trip() { + let fixture = Compact { + eviction_window: Percentage::new(0.2).unwrap(), + ..Compact::new() + }; + + let toml = toml_edit::ser::to_string_pretty(&fixture).unwrap(); + + assert!( + toml.contains("eviction_window = 0.2\n"), + "expected `eviction_window = 0.2` in TOML output, got:\n{toml}" + ); + } + + #[test] + fn test_f64_eviction_window_deserialize_round_trip() { + let fixture = Compact { + eviction_window: Percentage::new(0.2).unwrap(), + ..Compact::new() + }; + let config_fixture = ForgeConfig::default().compact(fixture.clone()); + + let toml = toml_edit::ser::to_string_pretty(&config_fixture).unwrap(); + + let actual = ConfigReader::default() + .read_defaults() + .read_toml(&toml) + .build() + .unwrap(); + let actual = actual.compact.expect("compact config should deserialize"); + + assert_eq!(actual.eviction_window, fixture.eviction_window); + } + + #[test] + fn test_token_threshold_percentage_round_trip() { + let fixture = Compact { + token_threshold_percentage: Some(Percentage::new(0.7).unwrap()), + ..Compact::new() + }; + let config_fixture = ForgeConfig::default().compact(fixture.clone()); + + let toml = toml_edit::ser::to_string_pretty(&config_fixture).unwrap(); + + assert!( + toml.contains("token_threshold_percentage = 0.7\n"), + "expected `token_threshold_percentage = 0.7` in TOML output, got:\n{toml}" + ); + + let actual = ConfigReader::default() + .read_defaults() + .read_toml(&toml) + .build() + .unwrap(); + let actual = actual.compact.expect("compact config should deserialize"); + + assert_eq!( + actual.token_threshold_percentage, + fixture.token_threshold_percentage + ); + } + + #[test] + fn test_token_threshold_percentage_rejects_out_of_range() { + let toml = "[compact]\ntoken_threshold_percentage = 1.5\n"; + + let result = ConfigReader::default() + .read_defaults() + .read_toml(toml) + .build(); + + assert!( + result.is_err(), + "expected error for token_threshold_percentage = 1.5, got: {:?}", + result.ok() + ); + } + + #[test] + fn test_eviction_window_rejects_out_of_range() { + let toml = "[compact]\neviction_window = 1.5\n"; + + let result = ConfigReader::default() + .read_defaults() + .read_toml(toml) + .build(); + + assert!( + result.is_err(), + "expected error for eviction_window = 1.5, got: {:?}", + result.ok() + ); + } + + #[test] + fn test_update_frequency_never_round_trip() { + let fixture = + ForgeConfig::default().updates(Update::default().frequency(UpdateFrequency::Never)); + + let toml = toml_edit::ser::to_string_pretty(&fixture).unwrap(); + + assert!( + toml.contains("frequency = \"never\"\n"), + "expected `frequency = \"never\"` in TOML output, got:\n{toml}" + ); + + let actual = ConfigReader::default() + .read_defaults() + .read_toml(&toml) + .build() + .unwrap(); + + let expected = Some( + Update::default() + .frequency(UpdateFrequency::Never) + .auto_update(true), + ); + assert_eq!(actual.updates, expected); + } +} diff --git a/crates/forge_config/src/config.rs b/crates/forge_config/src/config.rs new file mode 100644 index 0000000000000000000000000000000000000000..5c7ed51f90e451f5eaadca4d5ff88b44aeb95870 --- /dev/null +++ b/crates/forge_config/src/config.rs @@ -0,0 +1,516 @@ +use std::collections::HashMap; +use std::path::PathBuf; + +use derive_setters::Setters; +use fake::Dummy; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +use crate::reader::ConfigReader; +use crate::writer::ConfigWriter; +use crate::{ + AutoDumpFormat, Compact, Decimal, HttpConfig, ModelConfig, ReasoningConfig, RetryConfig, Update, +}; + +/// Wire protocol a provider uses for chat completions. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, Dummy)] +pub enum ProviderResponseType { + OpenAI, + OpenAIResponses, + Anthropic, + Bedrock, + Google, + OpenCode, +} + +/// Category of a provider. +#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize, JsonSchema, Dummy)] +#[serde(rename_all = "snake_case")] +pub enum ProviderTypeEntry { + /// LLM provider for chat completions. + #[default] + Llm, + /// Context engine provider for code indexing and search. + ContextEngine, +} + +/// Authentication method supported by a provider. +/// +/// Only the simple (non-OAuth) methods are available here; providers that +/// require OAuth device or authorization-code flows must be configured via the +/// file-based `provider.json` override instead. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, Dummy)] +#[serde(rename_all = "snake_case")] +pub enum ProviderAuthMethod { + ApiKey, + GoogleAdc, +} + +/// A URL parameter variable for a provider, used to substitute template +/// variables in URL strings. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, Dummy)] +#[serde(rename_all = "snake_case")] +pub struct ProviderUrlParam { + /// The environment variable name used as the template variable key. + pub name: String, + /// Optional preset values for this parameter shown as suggestions in the + /// UI. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub options: Vec, + /// Whether this parameter is optional. When `true`, the parameter may be + /// left blank without causing an error. + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub optional: bool, +} + +/// Source of models for a provider: either a URL to fetch them from or a +/// static list defined inline. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, Dummy)] +#[serde(untagged)] +pub enum ModelListConfig { + /// URL template used to fetch the model list dynamically. + Url(String), + /// A static list of models defined directly in the configuration. + Hardcoded(Vec), +} + +/// A single provider entry defined inline in `forge.toml`. +/// +/// Inline providers are merged with the built-in provider list; entries with +/// the same `id` override the corresponding built-in entry field-by-field, +/// while entries with a new `id` are appended to the list. +#[derive(Debug, Default, Clone, PartialEq, Serialize, Deserialize, JsonSchema, Dummy)] +#[serde(rename_all = "snake_case")] +pub struct ProviderEntry { + /// Unique provider identifier used in model paths (e.g. `"my_provider"`). + pub id: String, + /// Environment variable holding the API key for this provider. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub api_key_var: Option, + /// URL template for chat completions; may contain `{{VAR}}` placeholders + /// that are substituted from the credential's url params. + pub url: String, + /// Model source: either a URL template for dynamic discovery or a static + /// list of models defined inline. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub models: Option, + /// Wire protocol used by this provider. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub response_type: Option, + /// Environment variables whose values are substituted into `{{VAR}}` + /// placeholders in the `url` and `models` templates. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub url_param_vars: Vec, + /// Additional HTTP headers sent with every request to this provider. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub custom_headers: Option>, + /// Provider category; defaults to `llm` when omitted. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub provider_type: Option, + /// Authentication methods supported by this provider; defaults to + /// `["api_key"]` when omitted. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub auth_methods: Vec, +} + +/// Top-level Forge configuration merged from all sources (defaults, file, +/// environment). +#[derive(Default, Debug, Setters, Clone, PartialEq, Serialize, Deserialize, JsonSchema, Dummy)] +#[serde(rename_all = "snake_case")] +#[setters(strip_option)] +pub struct ForgeConfig { + /// Retry settings applied at the system level to all IO operations. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub retry: Option, + /// Maximum number of lines returned by a single file search operation. + #[serde(default)] + pub max_search_lines: usize, + /// Maximum number of bytes returned by a single file search operation. + #[serde(default)] + pub max_search_result_bytes: usize, + /// Maximum number of characters returned from a URL fetch. + #[serde(default)] + pub max_fetch_chars: usize, + /// Maximum number of lines captured from the leading portion of shell + /// command output. + #[serde(default)] + pub max_stdout_prefix_lines: usize, + /// Maximum number of lines captured from the trailing portion of shell + /// command output. + #[serde(default)] + pub max_stdout_suffix_lines: usize, + /// Maximum number of characters per line in shell command output. + #[serde(default)] + pub max_stdout_line_chars: usize, + /// Maximum number of characters per line when reading a file. + #[serde(default)] + pub max_line_chars: usize, + /// Maximum number of lines read from a file in a single operation. + #[serde(default)] + pub max_read_lines: u64, + /// Maximum number of files read in a single batch operation. + #[serde(default)] + pub max_file_read_batch_size: usize, + /// HTTP client settings including proxy, TLS, and timeout configuration. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub http: Option, + /// Maximum file size in bytes permitted for read operations. + #[serde(default)] + pub max_file_size_bytes: u64, + /// Maximum image file size in bytes permitted for read operations. + #[serde(default)] + pub max_image_size_bytes: u64, + /// Maximum time in seconds a single tool call may run before being + /// cancelled. + #[serde(default)] + pub tool_timeout_secs: u64, + /// Whether to automatically open HTML dump files in the browser after + /// creation. + #[serde(default)] + pub auto_open_dump: bool, + /// Directory where debug request files are written; disabled when absent. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub debug_requests: Option, + /// Path to the conversation history file; defaults to the global history + /// location when absent. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub custom_history_path: Option, + /// Maximum number of conversations shown in the conversation list. + #[serde(default)] + pub max_conversations: usize, + /// Maximum number of candidate results returned from the initial semantic + /// search vector query. + #[serde(default)] + pub max_sem_search_results: usize, + /// Number of top results retained after re-ranking in semantic search. + #[serde(default)] + pub sem_search_top_k: usize, + /// Base URL of the Forge services API used for semantic search and + /// indexing. + #[serde(default)] + #[dummy(expr = "\"https://api.forgecode.dev/api\".to_string()")] + pub services_url: String, + /// Maximum number of file extensions included in the agent system prompt. + #[serde(default)] + pub max_extensions: usize, + /// Format used when automatically creating a session dump after task + /// completion; disabled when absent. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub auto_dump: Option, + /// Maximum number of files read concurrently during batch operations. + #[serde(default)] + pub max_parallel_file_reads: usize, + /// Time-to-live in seconds for the cached model API list. + #[serde(default)] + pub model_cache_ttl_secs: u64, + /// Default model and provider configuration used when not overridden by + /// individual agents. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub session: Option, + /// Model and provider configuration used for commit message generation. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub commit: Option, + /// Whether `forge commit` should override `GIT_COMMITTER_NAME` and + /// `GIT_COMMITTER_EMAIL` with the Forge identity. Defaults to `true` via + /// the embedded `.forge.toml` defaults. + #[serde(default)] + pub use_forge_committer: bool, + /// Maximum number of recent commits included as context for commit message + /// generation. + #[serde(default)] + pub max_commit_count: usize, + /// Model and provider configuration used for shell command suggestion + /// generation. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub suggest: Option, + + // --- Workflow fields --- + /// Configuration for automatic Forge updates. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub updates: Option, + + /// Output randomness for all agents; lower values are deterministic, higher + /// values are creative (0.0–2.0). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub temperature: Option, + + /// Nucleus sampling threshold for all agents; limits token selection to the + /// top cumulative probability mass (0.0–1.0). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub top_p: Option, + + /// Top-k vocabulary cutoff for all agents; restricts sampling to the k + /// highest-probability tokens (1–1000). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub top_k: Option, + + /// Maximum tokens the model may generate per response for all agents + /// (1–100,000). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub max_tokens: Option, + + /// Maximum tool failures per turn before the orchestrator forces + /// completion. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub max_tool_failure_per_turn: Option, + + /// Maximum number of requests that can be made in a single turn. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub max_requests_per_turn: Option, + + /// Context compaction settings applied to all agents; falls back to each + /// agent's individual setting when absent. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub compact: Option, + + /// Whether restricted mode is active; when enabled, tool execution requires + /// explicit permission grants. + #[serde(default)] + pub restricted: bool, + + /// Whether tool use is supported in the current environment; when false, + /// all tool calls are disabled. + #[serde(default)] + pub tool_supported: bool, + + /// Reasoning configuration applied to all agents; controls effort level, + /// token budget, and visibility of the model's thinking process. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub reasoning: Option, + + /// Additional provider definitions merged with the built-in provider list. + /// + /// Entries with an `id` matching a built-in provider override its fields; + /// entries with a new `id` are appended and become available for model + /// selection. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub providers: Vec, + + /// Currency symbol displayed in the shell rprompt next to the session cost + /// (e.g. `"$"`, `"€"`, `"₹"`). Defaults to `"$"`. + #[serde(default)] + pub currency_symbol: String, + + /// Conversion rate applied to costs before display in the shell rprompt. + /// The raw USD cost is multiplied by this value, allowing costs to be shown + /// in a local currency. Defaults to `1.0` (no conversion). + #[serde(default)] + pub currency_conversion_rate: Decimal, + + /// Enables the pending todos hook that checks for incomplete todo items + /// when a task ends and reminds the LLM about them. + #[serde(default)] + pub verify_todos: bool, + + /// Switches patch replacement fallback from the legacy fuzzy-search range + /// lookup to the newer text-patch gRPC API. + /// Defaults to `false` so patching continues to use the legacy fallback + /// behavior unless explicitly enabled in `forge.toml`. + #[serde(default)] + pub use_text_patch_fallback: bool, + + /// Whether the deep research agent is available. + /// + /// When set to `true`, the Sage agent is added to the agent list and + /// the `:sage` app command is enabled. Defaults to `false`. + #[serde(default)] + pub research_subagent: bool, + + /// Enables subagent support via the task tool; when true the forge agent + /// gains access to the `task` tool for delegating work to specialised + /// sub-agents, and the `sage` research-only agent tool is removed. + /// When false the `task` tool is disabled and `sage` is available instead. + #[serde(default)] + pub subagents: bool, + + /// Enables automatic VS Code extension installation when Forge runs inside + /// VS Code and the extension is not already installed. + #[serde(default)] + pub auto_install_vscode_extension: bool, + + /// When `true`, all system messages in the conversation are merged into a + /// single leading system message before the request is sent. Enable this + /// for providers that reject requests containing system messages after + /// user or assistant turns (e.g. vLLM, NVIDIA NIM). + #[serde(default)] + pub merge_system_messages: bool, +} + +impl ForgeConfig { + /// Reads and merges configuration from all sources, returning the resolved + /// [`ForgeConfig`]. + /// + /// # Errors + /// + /// Returns an error if the config path cannot be resolved, the file cannot + /// be read, or deserialization fails. + pub fn read() -> crate::Result { + ConfigReader::default() + .read_legacy() + .read_defaults() + .read_global() + .read_env() + .build() + } + + /// Writes the configuration to the user config file. + /// + /// # Errors + /// + /// Returns an error if the configuration cannot be serialized or written to + /// disk. + pub fn write(&self) -> crate::Result<()> { + let path = ConfigReader::config_path(); + ConfigWriter::new(self.clone()).write(&path) + } +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::*; + use crate::reader::ConfigReader; + + #[test] + fn test_f32_temperature_round_trip() { + let fixture = ForgeConfig { temperature: Some(Decimal(0.1)), ..Default::default() }; + + let toml = toml_edit::ser::to_string_pretty(&fixture).unwrap(); + + assert!( + toml.contains("temperature = 0.1\n"), + "expected `temperature = 0.1` in TOML output, got:\n{toml}" + ); + } + + #[test] + fn test_f32_top_p_round_trip() { + let fixture = ForgeConfig { top_p: Some(Decimal(0.9)), ..Default::default() }; + + let toml = toml_edit::ser::to_string_pretty(&fixture).unwrap(); + + assert!( + toml.contains("top_p = 0.9\n"), + "expected `top_p = 0.9` in TOML output, got:\n{toml}" + ); + } + + #[test] + fn test_f32_temperature_deserialize_round_trip() { + let fixture = ForgeConfig { temperature: Some(Decimal(0.1)), ..Default::default() }; + + let toml = toml_edit::ser::to_string_pretty(&fixture).unwrap(); + + let actual = ConfigReader::default().read_toml(&toml).build().unwrap(); + + assert_eq!(actual.temperature, fixture.temperature); + } + + #[test] + fn test_provider_static_model_list_deserialization() { + let fixture = r#" +[[providers]] +id = "ollama" +url = "http://127.0.0.1:8000/v1/chat/completions" +response_type = "OpenAI" +auth_methods = ["api_key"] + +[[providers.models]] +id = "Qwen3.6-35B-A3b-q3-mlx" +name = "Qwen3.5-35B" +description = "Qwen local reasoning model with advanced problem-solving capabilities" +context_length = 262144 +tools_supported = true +supports_parallel_tool_calls = true +supports_reasoning = true +input_modalities = ["text"] + +[[providers.models]] +id = "llama3.2-3b" +name = "Llama 3.2 3B" +description = "Meta Llama 3.2 3B lightweight local model" +context_length = 131072 +tools_supported = true +supports_parallel_tool_calls = false +supports_reasoning = false +input_modalities = ["text"] +"#; + + let actual = ConfigReader::default().read_toml(fixture).build().unwrap(); + + let expected = vec![ProviderEntry { + id: "ollama".to_string(), + url: "http://127.0.0.1:8000/v1/chat/completions".to_string(), + response_type: Some(ProviderResponseType::OpenAI), + auth_methods: vec![ProviderAuthMethod::ApiKey], + models: Some(ModelListConfig::Hardcoded(vec![ + forge_domain::Model::new("Qwen3.6-35B-A3b-q3-mlx") + .name("Qwen3.5-35B".to_string()) + .description( + "Qwen local reasoning model with advanced problem-solving capabilities" + .to_string(), + ) + .context_length(262144) + .tools_supported(true) + .supports_parallel_tool_calls(true) + .supports_reasoning(true) + .input_modalities(vec![forge_domain::InputModality::Text]), + forge_domain::Model::new("llama3.2-3b") + .name("Llama 3.2 3B".to_string()) + .description("Meta Llama 3.2 3B lightweight local model".to_string()) + .context_length(131072) + .tools_supported(true) + .supports_parallel_tool_calls(false) + .supports_reasoning(false) + .input_modalities(vec![forge_domain::InputModality::Text]), + ])), + ..Default::default() + }]; + + assert_eq!(actual.providers, expected); + } + + #[test] + fn test_provider_url_model_list_deserialization() { + let fixture = r#" +[[providers]] +id = "my_provider" +url = "http://example.com/v1/chat/completions" +models = "http://example.com/v1/models" +"#; + + let actual = ConfigReader::default().read_toml(fixture).build().unwrap(); + + let expected = vec![ProviderEntry { + id: "my_provider".to_string(), + url: "http://example.com/v1/chat/completions".to_string(), + models: Some(ModelListConfig::Url( + "http://example.com/v1/models".to_string(), + )), + ..Default::default() + }]; + + assert_eq!(actual.providers, expected); + } + + #[test] + fn test_auto_install_vscode_extension_defaults_to_true() { + let actual = ConfigReader::default().read_defaults().build().unwrap(); + + assert_eq!(actual.auto_install_vscode_extension, true); + } + + #[test] + fn test_auto_install_vscode_extension_can_be_disabled() { + let toml = "auto_install_vscode_extension = false\n"; + + let actual = ConfigReader::default() + .read_defaults() + .read_toml(toml) + .build() + .unwrap(); + + assert_eq!(actual.auto_install_vscode_extension, false); + } +} diff --git a/crates/forge_config/src/decimal.rs b/crates/forge_config/src/decimal.rs new file mode 100644 index 0000000000000000000000000000000000000000..512714b804142f45105de8865858cde766d423ef --- /dev/null +++ b/crates/forge_config/src/decimal.rs @@ -0,0 +1,94 @@ +/// A floating-point newtype that serializes to two decimal places, preventing +/// `toml_edit` from emitting noisy bit-pattern approximations such as +/// `0.10000000149011612` or `0.20000000000000001`. +/// +/// The inner value is stored as `f64`. When used for fields that ultimately +/// require `f32`, callers should cast via `value() as f32`. +pub struct Decimal(pub f64); + +impl Decimal { + /// Returns the inner `f64` value. + pub fn value(&self) -> f64 { + self.0 + } +} + +impl std::fmt::Debug for Decimal { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + self.0.fmt(f) + } +} + +impl Clone for Decimal { + fn clone(&self) -> Self { + *self + } +} + +impl Copy for Decimal {} + +impl PartialEq for Decimal { + fn eq(&self, other: &Self) -> bool { + self.0 == other.0 + } +} + +impl PartialOrd for Decimal { + fn partial_cmp(&self, other: &Self) -> Option { + self.0.partial_cmp(&other.0) + } +} + +impl Default for Decimal { + fn default() -> Self { + Self(0.0) + } +} + +impl From for Decimal { + fn from(v: f64) -> Self { + Self(v) + } +} + +impl From for Decimal { + fn from(v: f32) -> Self { + Self(v as f64) + } +} + +impl From for f64 { + fn from(d: Decimal) -> Self { + d.0 + } +} + +impl schemars::JsonSchema for Decimal { + fn schema_name() -> std::borrow::Cow<'static, str> { + f64::schema_name() + } + + fn json_schema(r#gen: &mut schemars::generate::SchemaGenerator) -> schemars::Schema { + f64::json_schema(r#gen) + } +} + +impl fake::Dummy for Decimal { + fn dummy_with_rng(_: &fake::Faker, rng: &mut R) -> Self { + use fake::Fake; + Self((0.0f64..2.0f64).fake_with_rng(rng)) + } +} + +impl serde::Serialize for Decimal { + fn serialize(&self, serializer: S) -> Result { + let formatted: f64 = format!("{:.2}", self.0).parse().unwrap(); + serializer.serialize_f64(formatted) + } +} + +impl<'de> serde::Deserialize<'de> for Decimal { + fn deserialize>(deserializer: D) -> Result { + Ok(Self(f64::deserialize(deserializer)?)) + } +} diff --git a/crates/forge_config/src/error.rs b/crates/forge_config/src/error.rs new file mode 100644 index 0000000000000000000000000000000000000000..ed1b867b70e5f6e131b6b0450ee91eebe409ce3e --- /dev/null +++ b/crates/forge_config/src/error.rs @@ -0,0 +1,18 @@ +/// Errors produced by the `forge_config` crate. +#[derive(Debug, thiserror::Error)] +pub enum Error { + /// Failed to read or parse configuration from a file or environment. + #[error("Config error: {0}")] + Config(#[from] config::ConfigError), + + /// Failed to serialize or write configuration. + #[error("Serialization error: {0}")] + Serialization(#[from] toml_edit::ser::Error), + + /// An I/O error occurred while reading or writing configuration files. + #[error("IO error: {0}")] + Io(#[from] std::io::Error), + + #[error("JSON error: {0}")] + Serde(#[from] serde_json::Error), +} diff --git a/crates/forge_config/src/http.rs b/crates/forge_config/src/http.rs new file mode 100644 index 0000000000000000000000000000000000000000..25ae77c77fa6f76c9b3f6aa022eae6d763aee68c --- /dev/null +++ b/crates/forge_config/src/http.rs @@ -0,0 +1,95 @@ +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +/// TLS version enum for configuring TLS protocol versions. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, fake::Dummy)] +#[serde(rename_all = "snake_case")] +pub enum TlsVersion { + #[serde(rename = "1.0")] + V1_0, + #[serde(rename = "1.1")] + V1_1, + #[serde(rename = "1.2")] + V1_2, + #[serde(rename = "1.3")] + V1_3, +} + +/// TLS backend option. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, fake::Dummy)] +#[serde(rename_all = "snake_case")] +pub enum TlsBackend { + #[serde(rename = "default")] + Default, + #[serde(rename = "rustls")] + Rustls, +} + +/// HTTP client configuration. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, fake::Dummy)] +#[serde(rename_all = "snake_case")] +pub struct HttpConfig { + pub connect_timeout_secs: u64, + pub read_timeout_secs: u64, + pub pool_idle_timeout_secs: u64, + pub pool_max_idle_per_host: usize, + pub max_redirects: usize, + pub hickory: bool, + pub tls_backend: TlsBackend, + /// Minimum TLS protocol version to use + pub min_tls_version: Option, + /// Maximum TLS protocol version to use + pub max_tls_version: Option, + /// Adaptive window sizing for improved flow control + pub adaptive_window: bool, + /// Keep-alive interval in seconds + pub keep_alive_interval_secs: Option, + /// Keep-alive timeout in seconds + pub keep_alive_timeout_secs: u64, + /// Keep-alive while connection is idle + pub keep_alive_while_idle: bool, + /// Accept invalid certificates + pub accept_invalid_certs: bool, + /// Paths to root certificate files + pub root_cert_paths: Option>, +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::*; + + #[test] + fn test_http_config_fields() { + let config = HttpConfig { + connect_timeout_secs: 30, + read_timeout_secs: 900, + pool_idle_timeout_secs: 90, + pool_max_idle_per_host: 5, + max_redirects: 10, + hickory: false, + tls_backend: TlsBackend::Default, + min_tls_version: None, + max_tls_version: None, + adaptive_window: true, + keep_alive_interval_secs: Some(60), + keep_alive_timeout_secs: 10, + keep_alive_while_idle: true, + accept_invalid_certs: false, + root_cert_paths: None, + }; + assert_eq!(config.connect_timeout_secs, 30); + assert_eq!(config.adaptive_window, true); + } + + #[test] + fn test_tls_version_variants() { + assert_eq!(TlsVersion::V1_3, TlsVersion::V1_3); + } + + #[test] + fn test_tls_backend_variants() { + assert_eq!(TlsBackend::Default, TlsBackend::Default); + } +} diff --git a/crates/forge_config/src/legacy.rs b/crates/forge_config/src/legacy.rs new file mode 100644 index 0000000000000000000000000000000000000000..e8699ff9a32b573b9bf01551444ab32ad8339a32 --- /dev/null +++ b/crates/forge_config/src/legacy.rs @@ -0,0 +1,81 @@ +use std::collections::HashMap; +use std::path::PathBuf; + +use serde::Deserialize; + +use crate::{ForgeConfig, ModelConfig}; + +/// Intermediate representation of the legacy `~/forge/.config.json` format. +/// +/// This format stores the active provider as a top-level string and models as +/// a map from provider ID to model ID, which differs from the TOML config's +/// nested `session`, `commit`, and `suggest` sub-objects. +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct LegacyConfig { + /// The active provider ID (e.g. `"anthropic"`). + #[serde(default)] + provider: Option, + /// Map from provider ID to the model ID to use with that provider. + #[serde(default)] + model: HashMap, + /// Commit message generation provider/model pair. + #[serde(default)] + commit: Option, + /// Shell command suggestion provider/model pair. + #[serde(default)] + suggest: Option, +} + +/// A provider/model pair as expressed in the legacy JSON config. +#[derive(Debug, Deserialize)] +struct LegacyModelRef { + provider: Option, + model: Option, +} + +impl LegacyConfig { + /// Reads the legacy `~/forge/.config.json` file at `path`, parses it, and + /// returns the equivalent TOML representation as a [`String`]. + /// + /// Because every field in [`ForgeConfig`] is `Option`, fields not covered + /// by the legacy format are `None` and omitted from the serialized TOML, + /// so they cannot overwrite values from lower-priority config layers. + /// + /// # Errors + /// + /// Returns an error if the file cannot be read, the JSON is invalid, or the + /// resulting config cannot be serialized to TOML. + pub(crate) fn read(path: &PathBuf) -> crate::Result { + let contents = std::fs::read_to_string(path)?; + let config = serde_json::from_str::(&contents)?; + let forge_config = config.into_forge_config(); + let content = toml_edit::ser::to_string_pretty(&forge_config)?; + Ok(content) + } + + /// Converts a [`LegacyConfig`] into the fields of [`ForgeConfig`] that it + /// covers, leaving all other fields at their defaults (`None`). + fn into_forge_config(self) -> ForgeConfig { + let session = self.provider.as_deref().and_then(|provider_id| { + self.model + .get(provider_id) + .cloned() + .map(|model_id| ModelConfig { provider_id: provider_id.to_string(), model_id }) + }); + + let commit = self.commit.and_then(|c| { + c.provider + .zip(c.model) + .map(|(provider_id, model_id)| ModelConfig { provider_id, model_id }) + }); + + let suggest = self.suggest.and_then(|s| { + s.provider + .zip(s.model) + .map(|(provider_id, model_id)| ModelConfig { provider_id, model_id }) + }); + + ForgeConfig { session, commit, suggest, ..Default::default() } + } +} diff --git a/crates/forge_config/src/lib.rs b/crates/forge_config/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..cc253277e4d2a8c8260ed52be1f2c3b4ea30fdeb --- /dev/null +++ b/crates/forge_config/src/lib.rs @@ -0,0 +1,29 @@ +mod auto_dump; +mod compact; +mod config; +mod decimal; +mod error; +mod http; +mod legacy; +mod model; +mod percentage; +mod reader; +mod reasoning; +mod retry; +mod writer; + +pub use auto_dump::*; +pub use compact::*; +pub use config::*; +pub use decimal::*; +pub use error::Error; +pub use http::*; +pub use model::*; +pub use percentage::*; +pub use reader::*; +pub use reasoning::*; +pub use retry::*; +pub use writer::*; + +/// A `Result` type alias for this crate's [`Error`] type. +pub type Result = std::result::Result; diff --git a/crates/forge_config/src/model.rs b/crates/forge_config/src/model.rs new file mode 100644 index 0000000000000000000000000000000000000000..e48d8e7c4be9e2dad428445570b881b9b85be9e3 --- /dev/null +++ b/crates/forge_config/src/model.rs @@ -0,0 +1,26 @@ +use derive_setters::Setters; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +/// A type alias for a provider identifier string. +pub type ProviderId = String; + +/// A type alias for a model identifier string. +pub type ModelId = String; + +/// Pairs a provider and model together for a specific operation. +#[derive(Debug, Setters, Clone, PartialEq, Serialize, Deserialize, JsonSchema, fake::Dummy)] +#[setters(into)] +pub struct ModelConfig { + /// The provider to use for this operation. + pub provider_id: String, + /// The model to use for this operation. + pub model_id: String, +} + +impl ModelConfig { + /// Creates a new ModelConfig with the given provider and model IDs. + pub fn new(provider_id: impl Into, model_id: impl Into) -> Self { + Self { provider_id: provider_id.into(), model_id: model_id.into() } + } +} diff --git a/crates/forge_config/src/percentage.rs b/crates/forge_config/src/percentage.rs new file mode 100644 index 0000000000000000000000000000000000000000..ad299e3fc878adca965e399fb76edb0ba41b2414 --- /dev/null +++ b/crates/forge_config/src/percentage.rs @@ -0,0 +1,136 @@ +use crate::Decimal; + +/// A percentage value constrained to `[0.0, 1.0]`, built on top of [`Decimal`] +/// to inherit its two-decimal-place serialization and clean TOML output. +/// +/// Validation is enforced at deserialization time, so any config file with an +/// out-of-range value is rejected with a descriptive error. +#[derive(Default)] +pub struct Percentage(Decimal); + +impl Percentage { + const MIN: f64 = 0.0; + const MAX: f64 = 1.0; + + /// Construct a validated `Percentage`. + /// + /// # Errors + /// + /// Returns `Err` if `value` is outside `[0.0, 1.0]`. + pub fn new(value: f64) -> Result { + if (Self::MIN..=Self::MAX).contains(&value) { + Ok(Self(Decimal(value))) + } else { + Err(format!( + "value must be between {} and {}, got {value}", + Self::MIN, + Self::MAX + )) + } + } + + /// Returns the inner `f64` value. + pub fn value(&self) -> f64 { + self.0.value() + } +} + +impl std::fmt::Debug for Percentage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + self.0.fmt(f) + } +} + +impl Clone for Percentage { + fn clone(&self) -> Self { + *self + } +} + +impl Copy for Percentage {} + +impl PartialEq for Percentage { + fn eq(&self, other: &Self) -> bool { + self.0 == other.0 + } +} + +impl PartialOrd for Percentage { + fn partial_cmp(&self, other: &Self) -> Option { + self.0.partial_cmp(&other.0) + } +} + +impl From for Percentage { + fn from(v: f64) -> Self { + Self(Decimal::from(v)) + } +} + +impl From for f64 { + fn from(p: Percentage) -> Self { + p.0.into() + } +} + +impl schemars::JsonSchema for Percentage { + fn schema_name() -> std::borrow::Cow<'static, str> { + Decimal::schema_name() + } + + fn json_schema(r#gen: &mut schemars::generate::SchemaGenerator) -> schemars::Schema { + Decimal::json_schema(r#gen) + } +} + +impl fake::Dummy for Percentage { + fn dummy_with_rng(_: &fake::Faker, rng: &mut R) -> Self { + use fake::Fake; + Self(Decimal((0.0f64..=1.0f64).fake_with_rng::(rng))) + } +} + +impl serde::Serialize for Percentage { + fn serialize(&self, serializer: S) -> Result { + self.0.serialize(serializer) + } +} + +impl<'de> serde::Deserialize<'de> for Percentage { + fn deserialize>(deserializer: D) -> Result { + let decimal = Decimal::deserialize(deserializer)?; + Self::new(decimal.value()).map_err(serde::de::Error::custom) + } +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::*; + + #[test] + fn test_percentage_valid_range() { + assert!(Percentage::new(0.0).is_ok()); + assert!(Percentage::new(0.5).is_ok()); + assert!(Percentage::new(1.0).is_ok()); + } + + #[test] + fn test_percentage_rejects_out_of_range() { + assert!(Percentage::new(-0.1).is_err()); + assert!(Percentage::new(1.1).is_err()); + } + + #[test] + fn test_percentage_serializes_to_2dp() { + #[derive(serde::Serialize)] + struct Fixture { + value: Percentage, + } + let fixture = Fixture { value: Percentage::new(0.2).unwrap() }; + let actual = toml_edit::ser::to_string_pretty(&fixture).unwrap(); + let expected = "value = 0.2\n"; + assert_eq!(actual, expected); + } +} diff --git a/crates/forge_config/src/reader.rs b/crates/forge_config/src/reader.rs new file mode 100644 index 0000000000000000000000000000000000000000..ebcbedbfe83c3a425c035ca254376edf409c3074 --- /dev/null +++ b/crates/forge_config/src/reader.rs @@ -0,0 +1,314 @@ +use std::path::PathBuf; +use std::sync::LazyLock; + +use config::ConfigBuilder; +use config::builder::DefaultState; + +use crate::ForgeConfig; +use crate::legacy::LegacyConfig; + +/// Loads all `.env` files found while walking up from the current working +/// directory to the root, with priority given to closer (lower) directories. +/// Executed at most once per process. +static LOAD_DOT_ENV: LazyLock<()> = LazyLock::new(|| { + let cwd = std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")); + let mut paths = vec![]; + let mut current = PathBuf::new(); + + for component in cwd.components() { + current.push(component); + paths.push(current.clone()); + } + + paths.reverse(); + + for path in paths { + let env_file = path.join(".env"); + if env_file.is_file() { + dotenvy::from_path(&env_file).ok(); + } + } +}); + +/// Caches base-path resolution for the process lifetime. +static BASE_PATH: LazyLock = LazyLock::new(ConfigReader::resolve_base_path); + +/// Merges [`ForgeConfig`] from layered sources using a builder pattern. +#[derive(Default)] +pub struct ConfigReader { + builder: ConfigBuilder, +} + +impl ConfigReader { + /// Returns the path to the legacy JSON config file + /// (`~/.forge/.config.json`). + pub fn config_legacy_path() -> PathBuf { + Self::base_path().join(".config.json") + } + + /// Returns the path to the primary TOML config file + /// (`~/.forge/.forge.toml`). + pub fn config_path() -> PathBuf { + Self::base_path().join(".forge.toml") + } + + /// Returns the base directory for all Forge config files. + /// + /// Resolution order: + /// 1. `FORGE_CONFIG` environment variable, if set. + /// 2. `~/forge` (legacy path), if that directory exists, so users who have + /// not yet run `forge config migrate` continue to read from their + /// existing directory without disruption. + /// 3. `~/.forge` as the default path. + pub fn base_path() -> PathBuf { + BASE_PATH.clone() + } + + fn resolve_base_path() -> PathBuf { + if let Ok(path) = std::env::var("FORGE_CONFIG") { + return PathBuf::from(path); + } + + let base = dirs::home_dir().unwrap_or(PathBuf::from(".")); + let path = base.join("forge"); + + // Prefer ~/forge (legacy) when it exists so existing users are not + // disrupted; fall back to ~/.forge as the default. + if path.exists() { + tracing::info!("Using legacy path"); + return path; + } + + tracing::info!("Using new path"); + base.join(".forge") + } + + /// Adds the provided TOML string as a config source without touching the + /// filesystem. + pub fn read_toml(mut self, contents: &str) -> Self { + self.builder = self + .builder + .add_source(config::File::from_str(contents, config::FileFormat::Toml)); + + self + } + + /// Adds the embedded default config (`../.forge.toml`) as a source. + pub fn read_defaults(self) -> Self { + let defaults = include_str!("../.forge.toml"); + + self.read_toml(defaults) + } + + /// Adds `FORGE_`-prefixed environment variables as a config source. + pub fn read_env(mut self) -> Self { + self.builder = self.builder.add_source( + config::Environment::with_prefix("FORGE") + .prefix_separator("_") + .separator("__") + .try_parsing(true) + .list_separator(",") + .with_list_parse_key("retry.status_codes") + .with_list_parse_key("http.root_cert_paths"), + ); + + self + } + + /// Builds and deserializes all accumulated sources into a [`ForgeConfig`]. + /// + /// Triggers `.env` file loading (at most once per process) by walking up + /// the directory tree from the current working directory, with closer + /// directories taking priority. + /// + /// # Errors + /// + /// Returns an error if the configuration cannot be built or deserialized. + pub fn build(self) -> crate::Result { + *LOAD_DOT_ENV; + let config = self.builder.build()?; + Ok(config.try_deserialize::()?) + } + + /// Adds `~/.forge/.forge.toml` as a config source, silently skipping if + /// absent. + pub fn read_global(mut self) -> Self { + let path = Self::config_path(); + self.builder = self + .builder + .add_source(config::File::from(path).required(false)); + self + } + + /// Reads `~/.forge/.config.json` (legacy format) and adds it as a source, + /// silently skipping errors. + pub fn read_legacy(self) -> Self { + let content = LegacyConfig::read(&Self::config_legacy_path()); + if let Ok(content) = content { + self.read_toml(&content) + } else { + self + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::{Mutex, MutexGuard}; + + use pretty_assertions::assert_eq; + + use super::*; + use crate::ModelConfig; + + /// Serializes tests that mutate environment variables to prevent races. + static ENV_MUTEX: Mutex<()> = Mutex::new(()); + + /// Holds env vars set for a test's duration and removes them on drop, while + /// holding [`ENV_MUTEX`]. + struct EnvGuard { + keys: Vec<&'static str>, + _lock: MutexGuard<'static, ()>, + } + + impl EnvGuard { + /// Acquires [`ENV_MUTEX`], sets each `(key, value)` pair in the + /// environment, and removes each key in `remove` if present. All + /// set keys are cleaned up on drop. + #[must_use] + fn set(pairs: &[(&'static str, &str)]) -> Self { + Self::set_and_remove(pairs, &[]) + } + + /// Like [`set`] but also removes the listed keys before the test runs. + #[must_use] + fn set_and_remove(pairs: &[(&'static str, &str)], remove: &[&'static str]) -> Self { + let lock = ENV_MUTEX.lock().unwrap_or_else(|e| e.into_inner()); + let keys = pairs.iter().map(|(k, _)| *k).collect(); + for key in remove { + unsafe { std::env::remove_var(key) }; + } + for (key, value) in pairs { + unsafe { std::env::set_var(key, value) }; + } + Self { keys, _lock: lock } + } + } + + impl Drop for EnvGuard { + fn drop(&mut self) { + for key in &self.keys { + unsafe { std::env::remove_var(key) }; + } + } + } + + #[test] + fn test_base_path_uses_forge_config_env_var() { + let _guard = EnvGuard::set(&[("FORGE_CONFIG", "/custom/forge/dir")]); + let actual = ConfigReader::resolve_base_path(); + let expected = PathBuf::from("/custom/forge/dir"); + assert_eq!(actual, expected); + } + + #[test] + fn test_base_path_falls_back_to_home_dir_when_env_var_absent() { + // Hold the env mutex and ensure FORGE_CONFIG is absent so this test + // cannot race with test_base_path_uses_forge_config_env_var. + let _guard = EnvGuard::set_and_remove(&[], &["FORGE_CONFIG"]); + + let actual = ConfigReader::resolve_base_path(); + // Without FORGE_CONFIG set the path must be either "forge" (legacy, + // preferred when ~/forge exists) or ".forge" (default new path). + let name = actual.file_name().unwrap(); + assert!( + name == "forge" || name == ".forge", + "Expected base_path to end with 'forge' or '.forge', got: {:?}", + name + ); + } + + #[test] + fn test_read_parses_without_error() { + let actual = ConfigReader::default().read_defaults().build(); + assert!(actual.is_ok(), "read() failed: {:?}", actual.err()); + } + + #[test] + fn test_legacy_layer_does_not_overwrite_defaults() { + // Simulate what `read_legacy` does: serialize a ForgeConfig that only + // carries session/commit/suggest (all other fields are None) and layer + // it on top of the embedded defaults. The default values must survive. + let legacy = ForgeConfig { + session: Some(ModelConfig { + provider_id: "anthropic".to_string(), + model_id: "claude-3".to_string(), + }), + ..Default::default() + }; + let legacy_toml = toml_edit::ser::to_string_pretty(&legacy).unwrap(); + + let actual = ConfigReader::default() + // Read legacy first and then defaults + .read_toml(&legacy_toml) + .read_defaults() + .build() + .unwrap(); + + // Session should come from the legacy layer + assert_eq!( + actual.session, + Some(ModelConfig { + provider_id: "anthropic".to_string(), + model_id: "claude-3".to_string(), + }) + ); + + // Default values from .forge.toml must be retained, not reset to zero + assert_eq!(actual.max_parallel_file_reads, 64); + assert_eq!(actual.max_read_lines, 2000); + assert_eq!(actual.tool_timeout_secs, 300); + assert_eq!(actual.max_search_lines, 1000); + assert_eq!(actual.tool_supported, true); + } + + #[test] + fn test_read_session_from_env_vars() { + let _guard = EnvGuard::set(&[ + ("FORGE_SESSION__PROVIDER_ID", "fake-provider"), + ("FORGE_SESSION__MODEL_ID", "fake-model"), + ]); + + let actual = ConfigReader::default() + .read_defaults() + .read_env() + .build() + .unwrap(); + + let expected = Some(ModelConfig { + provider_id: "fake-provider".to_string(), + model_id: "fake-model".to_string(), + }); + assert_eq!(actual.session, expected); + } + + #[test] + fn test_use_forge_committer_defaults_to_true() { + let actual = ConfigReader::default().read_defaults().build().unwrap(); + + assert_eq!(actual.use_forge_committer, true); + } + + #[test] + fn test_use_forge_committer_can_be_disabled() { + let toml = "use_forge_committer = false\n"; + + let actual = ConfigReader::default() + .read_defaults() + .read_toml(toml) + .build() + .unwrap(); + + assert_eq!(actual.use_forge_committer, false); + } +} diff --git a/crates/forge_config/src/reasoning.rs b/crates/forge_config/src/reasoning.rs new file mode 100644 index 0000000000000000000000000000000000000000..2542d89dc64f859b2aa32e8d6d248cb2c0dc7d15 --- /dev/null +++ b/crates/forge_config/src/reasoning.rs @@ -0,0 +1,57 @@ +use derive_setters::Setters; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; +use strum_macros::Display as StrumDisplay; + +/// Controls the reasoning behaviour of a model, including effort level, token +/// budget, and visibility of the thinking process. +#[derive( + Default, Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, fake::Dummy, Setters, +)] +#[serde(rename_all = "snake_case")] +#[setters(strip_option)] +pub struct ReasoningConfig { + /// Controls the effort level of the model's reasoning. + /// Supported by openrouter and forge provider. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub effort: Option, + + /// Controls how many tokens the model can spend thinking. + /// Should be greater than 1024 but less than the overall max_tokens. + /// Supported by openrouter, anthropic, and forge provider. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub max_tokens: Option, + + /// When true, the model thinks deeply but the reasoning is hidden from the + /// caller. Supported by openrouter and forge provider. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub exclude: Option, + + /// Enables reasoning at the "medium" effort level with no exclusions. + /// Supported by openrouter, anthropic, and forge provider. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub enabled: Option, +} + +/// Effort level for model reasoning. +#[derive( + Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, fake::Dummy, StrumDisplay, +)] +#[serde(rename_all = "lowercase")] +#[strum(serialize_all = "lowercase")] +pub enum Effort { + /// No reasoning; skips the thinking step entirely. + None, + /// Minimal reasoning; fastest and cheapest. + Minimal, + /// Low reasoning effort. + Low, + /// Medium reasoning effort; the default for most providers. + Medium, + /// High reasoning effort. + High, + /// Extra-high reasoning effort (OpenAI / OpenRouter). + XHigh, + /// Maximum reasoning effort; only available on select Anthropic models. + Max, +} diff --git a/crates/forge_config/src/retry.rs b/crates/forge_config/src/retry.rs new file mode 100644 index 0000000000000000000000000000000000000000..cc0d917be58cb0a9dd2c1a8e412669807052b3ed --- /dev/null +++ b/crates/forge_config/src/retry.rs @@ -0,0 +1,48 @@ +use derive_setters::Setters; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +/// Configuration for retry mechanism. +#[derive( + Default, Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, fake::Dummy, Setters, +)] +#[serde(rename_all = "snake_case")] +#[setters(into)] +pub struct RetryConfig { + /// Initial backoff delay in milliseconds for retry operations + pub initial_backoff_ms: u64, + /// Minimum delay in milliseconds between retry attempts + pub min_delay_ms: u64, + /// Backoff multiplication factor for each retry attempt + pub backoff_factor: u64, + /// Maximum number of retry attempts + pub max_attempts: usize, + /// HTTP status codes that should trigger retries + pub status_codes: Vec, + /// Maximum delay between retries in seconds + pub max_delay_secs: Option, + /// Whether to suppress retry error logging and events + pub suppress_errors: bool, +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::*; + + #[test] + fn test_retry_config_fields() { + let config = RetryConfig { + initial_backoff_ms: 200, + min_delay_ms: 1000, + backoff_factor: 2, + max_attempts: 8, + status_codes: vec![429, 500, 502, 503, 504, 408, 522, 524, 520, 529], + max_delay_secs: None, + suppress_errors: false, + }; + assert_eq!(config.initial_backoff_ms, 200); + assert_eq!(config.suppress_errors, false); + } +} diff --git a/crates/forge_config/src/writer.rs b/crates/forge_config/src/writer.rs new file mode 100644 index 0000000000000000000000000000000000000000..14b70a6884568f85dee67a3a5b4cf98d0e8af64b --- /dev/null +++ b/crates/forge_config/src/writer.rs @@ -0,0 +1,40 @@ +use std::path::Path; + +use crate::ForgeConfig; + +/// Writes a [`ForgeConfig`] to the user configuration file on disk. +pub struct ConfigWriter { + config: ForgeConfig, +} + +impl ConfigWriter { + /// Creates a new `ConfigWriter` for the given configuration. + pub fn new(config: ForgeConfig) -> Self { + Self { config } + } + + /// Serializes and writes the configuration to `path`, creating all parent + /// directories recursively if they do not already exist. + /// + /// The output includes a leading `$schema` key pointing to the Forge + /// configuration JSON schema, which enables editor validation and + /// auto-complete. + /// + /// # Errors + /// + /// Returns an error if the configuration cannot be serialized or the file + /// cannot be written. + pub fn write(&self, path: &Path) -> crate::Result<()> { + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + + let config_toml = toml_edit::ser::to_string_pretty(&self.config)?; + let contents = + format!("\"$schema\" = \"https://forgecode.dev/schema.json\"\n\n{config_toml}"); + + std::fs::write(path, contents)?; + + Ok(()) + } +} diff --git a/crates/forge_config/tests/schema.rs b/crates/forge_config/tests/schema.rs new file mode 100644 index 0000000000000000000000000000000000000000..1588919ef618a7a1dd453ae2b5729fa9d2794a90 --- /dev/null +++ b/crates/forge_config/tests/schema.rs @@ -0,0 +1,30 @@ +use std::path::Path; + +use forge_config::ForgeConfig; +use pretty_assertions::assert_eq; + +#[tokio::test] +async fn generate_workflow_schema() -> anyhow::Result<()> { + let schema = schemars::schema_for!(ForgeConfig); + let generated_schema = serde_json::to_string_pretty(&schema)?; + + // Use the crate root directory for the schema file + let crate_root = env!("CARGO_MANIFEST_DIR"); + let schema_path = Path::new(crate_root).join("../../forge.schema.json"); + + if is_ci::uncached() { + // On CI: validate that the generated schema matches the committed file + let existing_schema = tokio::fs::read_to_string(&schema_path).await?; + assert_eq!( + generated_schema.trim(), + existing_schema.trim(), + "Generated workflow schema does not match the committed schema file. \ + Please run the test locally to update the schema file." + ); + } else { + // Locally: generate and write the schema file + tokio::fs::write(&schema_path, generated_schema).await?; + } + + Ok(()) +} diff --git a/crates/forge_embed/src/lib.rs b/crates/forge_embed/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..5dba27544984a0d03faf670b7b8097b8c09caf9b --- /dev/null +++ b/crates/forge_embed/src/lib.rs @@ -0,0 +1,39 @@ +use handlebars::Handlebars; +use include_dir::{Dir, DirEntry, File}; + +/// Returns an iterator over all files embedded in `dir`, recursively +/// descending into subdirectories. +pub fn files(dir: &'static Dir<'static>) -> impl Iterator> { + dir.entries().iter().flat_map(walk_entry) +} + +fn walk_entry( + entry: &'static DirEntry<'static>, +) -> Box>> { + match entry { + DirEntry::File(f) => Box::new(std::iter::once(f)), + DirEntry::Dir(d) => Box::new(d.entries().iter().flat_map(walk_entry)), + } +} + +/// Registers all files in `dir` (recursively) as Handlebars templates. +/// +/// Template names match the relative file path as returned by +/// [`File::path`] (e.g. `forge-system-prompt.md`). Panics if any file path +/// is not valid UTF-8, if any file content is not valid UTF-8, or if template +/// parsing fails. +pub fn register_templates(hb: &mut Handlebars<'_>, dir: &'static Dir<'static>) { + for file in files(dir) { + let name = file.path().to_str().unwrap_or_else(|| { + panic!( + "embedded template path '{:?}' is not valid UTF-8", + file.path() + ) + }); + let content = file + .contents_utf8() + .unwrap_or_else(|| panic!("embedded template '{}' is not valid UTF-8", name)); + hb.register_template_string(name, content) + .unwrap_or_else(|e| panic!("failed to register template '{}': {}", name, e)); + } +} diff --git a/crates/forge_eventsource/src/error.rs b/crates/forge_eventsource/src/error.rs new file mode 100644 index 0000000000000000000000000000000000000000..97576ffbabb08eaab426ba84ec08b11aa611ff6b --- /dev/null +++ b/crates/forge_eventsource/src/error.rs @@ -0,0 +1,59 @@ +use core::fmt; +use std::string::FromUtf8Error; + +use forge_eventsource_stream::EventStreamError; +use nom::error::Error as NomError; +#[cfg(doc)] +use reqwest::RequestBuilder; +use reqwest::header::HeaderValue; +use reqwest::{Error as ReqwestError, Response, StatusCode}; + +/// Error raised when a [`RequestBuilder`] cannot be cloned. See +/// [`RequestBuilder::try_clone`] for more information +#[derive(Debug, Clone, Copy)] +pub struct CannotCloneRequestError; + +impl fmt::Display for CannotCloneRequestError { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.write_str("expected a cloneable request") + } +} + +impl std::error::Error for CannotCloneRequestError {} + +/// Error raised by the EventSource stream fetching and parsing +#[derive(Debug, Error)] +pub enum Error { + /// Source stream is not valid UTF8 + #[error(transparent)] + Utf8(FromUtf8Error), + /// Source stream is not a valid EventStream + #[error(transparent)] + Parser(NomError), + /// The HTTP Request could not be completed + #[error(transparent)] + Transport(ReqwestError), + /// The `Content-Type` returned by the server is invalid + #[error("Invalid header value: {0:?}")] + InvalidContentType(HeaderValue, Box), + /// The status code returned by the server is invalid + #[error("Invalid status code: {0}")] + InvalidStatusCode(StatusCode, Box), + /// The `Last-Event-ID` cannot be formed into a Header to be submitted to + /// the server + #[error("Invalid `Last-Event-ID`: {0}")] + InvalidLastEventId(String), + /// The stream ended + #[error("Stream ended")] + StreamEnded, +} + +impl From> for Error { + fn from(err: EventStreamError) -> Self { + match err { + EventStreamError::Utf8(err) => Self::Utf8(err), + EventStreamError::Parser(err) => Self::Parser(err), + EventStreamError::Transport(err) => Self::Transport(err), + } + } +} diff --git a/crates/forge_eventsource/src/event_source.rs b/crates/forge_eventsource/src/event_source.rs new file mode 100644 index 0000000000000000000000000000000000000000..17b15262e57a04a13acd9afb73b201f8ca6a8930 --- /dev/null +++ b/crates/forge_eventsource/src/event_source.rs @@ -0,0 +1,290 @@ +use core::pin::Pin; +use std::time::Duration; + +use forge_eventsource_stream::Eventsource; +pub use forge_eventsource_stream::{Event as MessageEvent, EventStreamError}; +#[cfg(not(target_arch = "wasm32"))] +use futures_core::future::BoxFuture; +use futures_core::future::Future; +#[cfg(target_arch = "wasm32")] +use futures_core::future::LocalBoxFuture; +#[cfg(not(target_arch = "wasm32"))] +use futures_core::stream::BoxStream; +#[cfg(target_arch = "wasm32")] +use futures_core::stream::LocalBoxStream; +use futures_core::stream::Stream; +use futures_core::task::{Context, Poll}; +use futures_timer::Delay; +use pin_project_lite::pin_project; +use reqwest::header::{HeaderName, HeaderValue}; +use reqwest::{Error as ReqwestError, IntoUrl, RequestBuilder, Response, StatusCode}; + +use crate::error::{CannotCloneRequestError, Error}; +use crate::retry::{DEFAULT_RETRY, RetryPolicy}; + +#[cfg(not(target_arch = "wasm32"))] +type ResponseFuture = BoxFuture<'static, Result>; +#[cfg(target_arch = "wasm32")] +type ResponseFuture = LocalBoxFuture<'static, Result>; + +#[cfg(not(target_arch = "wasm32"))] +type EventStream = BoxStream<'static, Result>>; +#[cfg(target_arch = "wasm32")] +type EventStream = LocalBoxStream<'static, Result>>; + +type BoxedRetry = Box; + +/// The ready state of an [`EventSource`] +#[derive(Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd)] +#[repr(u8)] +pub enum ReadyState { + /// The EventSource is waiting on a response from the endpoint + Connecting = 0, + /// The EventSource is connected + Open = 1, + /// The EventSource is closed and no longer emitting Events + Closed = 2, +} + +pin_project! { +/// Provides the [`Stream`] implementation for the [`Event`] items. This wraps the +/// [`RequestBuilder`] and retries requests when they fail. +#[project = EventSourceProjection] +pub struct EventSource { + builder: RequestBuilder, + #[pin] + next_response: Option, + #[pin] + cur_stream: Option, + #[pin] + delay: Option, + is_closed: bool, + retry_policy: BoxedRetry, + last_event_id: String, + last_retry: Option<(usize, Duration)> +} +} + +impl EventSource { + /// Wrap a [`RequestBuilder`] + pub fn new(builder: RequestBuilder) -> Result { + let builder = builder.header( + reqwest::header::ACCEPT, + HeaderValue::from_static("text/event-stream"), + ); + let res_future = Box::pin(builder.try_clone().ok_or(CannotCloneRequestError)?.send()); + Ok(Self { + builder, + next_response: Some(res_future), + cur_stream: None, + delay: None, + is_closed: false, + retry_policy: Box::new(DEFAULT_RETRY), + last_event_id: String::new(), + last_retry: None, + }) + } + + /// Create a simple EventSource based on a GET request + pub fn get(url: T) -> Self { + Self::new(reqwest::Client::new().get(url)).unwrap() + } + + /// Close the EventSource stream and stop trying to reconnect + pub fn close(&mut self) { + self.is_closed = true; + } + + /// Set the retry policy + pub fn set_retry_policy(&mut self, policy: BoxedRetry) { + self.retry_policy = policy + } + + /// Get the last event id + pub fn last_event_id(&self) -> &str { + &self.last_event_id + } + + /// Get the current ready state + pub fn ready_state(&self) -> ReadyState { + if self.is_closed { + ReadyState::Closed + } else if self.delay.is_some() || self.next_response.is_some() { + ReadyState::Connecting + } else { + ReadyState::Open + } + } +} + +fn check_response(response: Response) -> Result { + match response.status() { + StatusCode::OK => {} + status => { + return Err(Error::InvalidStatusCode(status, Box::new(response))); + } + } + let content_type = + if let Some(content_type) = response.headers().get(&reqwest::header::CONTENT_TYPE) { + content_type + } else { + return Err(Error::InvalidContentType( + HeaderValue::from_static(""), + Box::new(response), + )); + }; + if content_type + .to_str() + .map_err(|_| ()) + .and_then(|s| s.parse::().map_err(|_| ())) + .map(|mime_type| { + matches!( + (mime_type.type_(), mime_type.subtype()), + (mime::TEXT, mime::EVENT_STREAM) + ) + }) + .unwrap_or(false) + { + Ok(response) + } else { + Err(Error::InvalidContentType( + content_type.clone(), + Box::new(response), + )) + } +} + +impl<'a> EventSourceProjection<'a> { + fn clear_fetch(&mut self) { + self.next_response.take(); + self.cur_stream.take(); + } + + fn retry_fetch(&mut self) -> Result<(), Error> { + self.cur_stream.take(); + let req = self.builder.try_clone().unwrap().header( + HeaderName::from_static("last-event-id"), + HeaderValue::from_str(self.last_event_id) + .map_err(|_| Error::InvalidLastEventId(self.last_event_id.clone()))?, + ); + let res_future = Box::pin(req.send()); + self.next_response.replace(res_future); + Ok(()) + } + + fn handle_response(&mut self, res: Response) { + self.last_retry.take(); + let mut stream = res.bytes_stream().eventsource(); + stream.set_last_event_id(self.last_event_id.clone()); + self.cur_stream.replace(Box::pin(stream)); + } + + fn handle_event(&mut self, event: &MessageEvent) { + *self.last_event_id = event.id.clone(); + if let Some(duration) = event.retry { + self.retry_policy.set_reconnection_time(duration) + } + } + + fn handle_error(&mut self, error: &Error) { + self.clear_fetch(); + if let Some(retry_delay) = self.retry_policy.retry(error, *self.last_retry) { + let retry_num = self.last_retry.map(|retry| retry.0).unwrap_or(1); + *self.last_retry = Some((retry_num, retry_delay)); + self.delay.replace(Delay::new(retry_delay)); + } else { + *self.is_closed = true; + } + } +} + +/// Events created by the [`EventSource`] +#[derive(Debug, Clone, Eq, PartialEq)] +pub enum Event { + /// The event fired when the connection is opened + Open, + /// The event fired when a [`MessageEvent`] is received + Message(MessageEvent), +} + +impl From for Event { + fn from(event: MessageEvent) -> Self { + Event::Message(event) + } +} + +impl Stream for EventSource { + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll> { + let mut this = self.project(); + + if *this.is_closed { + return Poll::Ready(None); + } + + if let Some(delay) = this.delay.as_mut().as_pin_mut() { + match delay.poll(cx) { + Poll::Ready(_) => { + this.delay.take(); + if let Err(err) = this.retry_fetch() { + *this.is_closed = true; + return Poll::Ready(Some(Err(err))); + } + } + Poll::Pending => return Poll::Pending, + } + } + + if let Some(response_future) = this.next_response.as_mut().as_pin_mut() { + match response_future.poll(cx) { + Poll::Ready(Ok(res)) => { + this.clear_fetch(); + match check_response(res) { + Ok(res) => { + this.handle_response(res); + return Poll::Ready(Some(Ok(Event::Open))); + } + Err(err) => { + *this.is_closed = true; + return Poll::Ready(Some(Err(err))); + } + } + } + Poll::Ready(Err(err)) => { + let err = Error::Transport(err); + this.handle_error(&err); + return Poll::Ready(Some(Err(err))); + } + Poll::Pending => { + return Poll::Pending; + } + } + } + + match this + .cur_stream + .as_mut() + .as_pin_mut() + .unwrap() + .as_mut() + .poll_next(cx) + { + Poll::Ready(Some(Err(err))) => { + let err = err.into(); + this.handle_error(&err); + Poll::Ready(Some(Err(err))) + } + Poll::Ready(Some(Ok(event))) => { + this.handle_event(&event); + Poll::Ready(Some(Ok(event.into()))) + } + Poll::Ready(None) => { + let err = Error::StreamEnded; + this.handle_error(&err); + Poll::Ready(Some(Err(err))) + } + Poll::Pending => Poll::Pending, + } + } +} diff --git a/crates/forge_eventsource/src/lib.rs b/crates/forge_eventsource/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..da303ea3901cbfc8f388be8d4a81a8886934bfab --- /dev/null +++ b/crates/forge_eventsource/src/lib.rs @@ -0,0 +1,33 @@ +//! Provides a simple wrapper for [`reqwest`] to provide an Event Source +//! implementation. You can learn more about Server Sent Events (SSE) take a +//! look at [the MDN docs](https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events) +//! This crate uses [`forge_eventsource_stream`] to wrap the underlying Bytes +//! stream, and retries failed requests. +//! +//! # Example +//! +//! ```ignore +//! let mut es = EventSource::get("http://localhost:8000/events"); +//! while let Some(event) = es.next().await { +//! match event { +//! Ok(Event::Open) => println!("Connection Open!"), +//! Ok(Event::Message(message)) => println!("Message: {:#?}", message), +//! Err(err) => { +//! println!("Error: {}", err); +//! es.close(); +//! } +//! } +//! } +//! ``` + +#[macro_use] +extern crate thiserror; + +mod error; +mod event_source; +mod reqwest_ext; +pub mod retry; + +pub use error::{CannotCloneRequestError, Error}; +pub use event_source::{Event, EventSource, ReadyState}; +pub use reqwest_ext::RequestBuilderExt; diff --git a/crates/forge_eventsource/src/reqwest_ext.rs b/crates/forge_eventsource/src/reqwest_ext.rs new file mode 100644 index 0000000000000000000000000000000000000000..d5562478d598cbd9a717fcda3d5b6be33d33d2af --- /dev/null +++ b/crates/forge_eventsource/src/reqwest_ext.rs @@ -0,0 +1,17 @@ +use reqwest::RequestBuilder; + +use crate::error::CannotCloneRequestError; +use crate::event_source::EventSource; + +/// Provides an easy interface to build an [`EventSource`] from a +/// [`RequestBuilder`] +pub trait RequestBuilderExt { + /// Create a new [`EventSource`] from a [`RequestBuilder`] + fn eventsource(self) -> Result; +} + +impl RequestBuilderExt for RequestBuilder { + fn eventsource(self) -> Result { + EventSource::new(self) + } +} diff --git a/crates/forge_eventsource/src/retry.rs b/crates/forge_eventsource/src/retry.rs new file mode 100644 index 0000000000000000000000000000000000000000..ed7f9b5b48c611b39c5424106a6d2a0b9e083c0b --- /dev/null +++ b/crates/forge_eventsource/src/retry.rs @@ -0,0 +1,120 @@ +//! Helpers to handle connection delays when receiving errors + +use std::time::Duration; + +use crate::error::Error; +#[cfg(doc)] +use crate::event_source::{Event, EventSource}; + +/// Describes how an [`EventSource`] should retry on receiving an [`enum@Error`] +pub trait RetryPolicy { + /// Submit a new retry delay based on the [`enum@Error`], last retry number + /// and duration, if available. A policy may also return `None` if it + /// does not want to retry + fn retry(&self, error: &Error, last_retry: Option<(usize, Duration)>) -> Option; + + /// Set a new reconnection time if received from an [`Event`] + fn set_reconnection_time(&mut self, duration: Duration); +} + +/// A [`RetryPolicy`] which backs off exponentially +#[derive(Debug, Clone)] +pub struct ExponentialBackoff { + /// The start of the backoff + pub start: Duration, + /// The factor of which to backoff by + pub factor: f64, + /// The maximum duration to delay + pub max_duration: Option, + /// The maximum number of retries before giving up + pub max_retries: Option, +} + +impl ExponentialBackoff { + /// Create a new exponential backoff retry policy + pub const fn new( + start: Duration, + factor: f64, + max_duration: Option, + max_retries: Option, + ) -> Self { + Self { start, factor, max_duration, max_retries } + } +} + +impl RetryPolicy for ExponentialBackoff { + fn retry(&self, _error: &Error, last_retry: Option<(usize, Duration)>) -> Option { + if let Some((retry_num, last_duration)) = last_retry { + if self.max_retries.is_none() || retry_num < self.max_retries.unwrap() { + let duration = last_duration.mul_f64(self.factor); + if let Some(max_duration) = self.max_duration { + Some(duration.min(max_duration)) + } else { + Some(duration) + } + } else { + None + } + } else { + Some(self.start) + } + } + fn set_reconnection_time(&mut self, duration: Duration) { + self.start = duration; + if let Some(max_duration) = self.max_duration { + self.max_duration = Some(max_duration.max(duration)) + } + } +} + +/// A [`RetryPolicy`] which always emits the same delay +#[derive(Debug, Clone)] +pub struct Constant { + /// The delay to return + pub delay: Duration, + /// The maximum number of retries to return before giving up + pub max_retries: Option, +} + +impl Constant { + /// Create a new constant retry policy + pub const fn new(delay: Duration, max_retries: Option) -> Self { + Self { delay, max_retries } + } +} + +impl RetryPolicy for Constant { + fn retry(&self, _error: &Error, last_retry: Option<(usize, Duration)>) -> Option { + if let Some((retry_num, _)) = last_retry { + if self.max_retries.is_none() || retry_num < self.max_retries.unwrap() { + Some(self.delay) + } else { + None + } + } else { + Some(self.delay) + } + } + fn set_reconnection_time(&mut self, duration: Duration) { + self.delay = duration; + } +} + +/// A [`RetryPolicy`] which never retries +#[derive(Debug, Clone, Copy, Default)] +pub struct Never; + +impl RetryPolicy for Never { + fn retry(&self, _error: &Error, _last_retry: Option<(usize, Duration)>) -> Option { + None + } + fn set_reconnection_time(&mut self, _duration: Duration) {} +} + +/// The default [`RetryPolicy`] when initializing an [`EventSource`] +pub const DEFAULT_RETRY: ExponentialBackoff = ExponentialBackoff::new( + Duration::from_millis(300), + 2., + Some(Duration::from_secs(5)), + None, +); diff --git a/crates/forge_infra/src/auth/error.rs b/crates/forge_infra/src/auth/error.rs new file mode 100644 index 0000000000000000000000000000000000000000..56f34d6e0e3f939bcea7d450b3c387d2eebe32c5 --- /dev/null +++ b/crates/forge_infra/src/auth/error.rs @@ -0,0 +1,37 @@ +use std::time::Duration; + +/// Authentication flow errors. +#[derive(Debug, Clone, thiserror::Error)] +pub enum Error { + /// Authentication initiation failed. + #[error("Authentication initiation failed: {0}")] + InitiationFailed(String), + + /// Authentication timed out waiting for user. + #[error("Authentication timed out after {0:?}")] + Timeout(Duration), + + /// Device code or session expired before completion. + #[error("Device code or session expired")] + Expired, + + /// User denied authorization request. + #[error("User denied authorization")] + Denied, + + /// Polling operation failed due to network or server error. + #[error("Polling failed: {0}")] + PollFailed(String), + + /// Authentication completion failed. + #[error("Authentication completion failed: {0}")] + CompletionFailed(String), + + /// Token refresh operation failed. + #[error("Token refresh failed: {0}")] + RefreshFailed(String), + + /// Invalid authentication context for the flow type. + #[error("Invalid context: {0}")] + InvalidContext(String), +} diff --git a/crates/forge_infra/src/auth/http/anthropic.rs b/crates/forge_infra/src/auth/http/anthropic.rs new file mode 100644 index 0000000000000000000000000000000000000000..85f77d53eb639848b163b0a0384e7edf5fb4da69 --- /dev/null +++ b/crates/forge_infra/src/auth/http/anthropic.rs @@ -0,0 +1,148 @@ +use forge_app::OAuthHttpProvider; +use forge_domain::{AuthCodeParams, OAuthConfig, OAuthTokenResponse}; +use oauth2::PkceCodeChallenge; +use serde::Serialize; + +use crate::auth::util::build_http_client; + +/// Anthropic Provider - Non-standard PKCE implementation +/// Quirk: state parameter equals PKCE verifier +#[allow(unused)] +pub struct AnthropicHttpProvider; + +#[allow(unused)] +#[derive(Debug, Serialize)] +struct AnthropicTokenRequest { + /// Authorization code from callback + code: String, + /// State parameter (equals PKCE verifier) + state: String, + /// Must be "authorization_code" + grant_type: String, + /// OAuth client ID + client_id: String, + /// Redirect URI (must match authorization request) + #[serde(skip_serializing_if = "Option::is_none")] + redirect_uri: Option, + /// PKCE code verifier + code_verifier: String, +} + +#[async_trait::async_trait] +impl OAuthHttpProvider for AnthropicHttpProvider { + async fn build_auth_url(&self, config: &OAuthConfig) -> anyhow::Result { + // Anthropic quirk: state = verifier + let (challenge, verifier) = PkceCodeChallenge::new_random_sha256(); + + let mut url = config.auth_url.clone(); + url.query_pairs_mut() + .append_pair("client_id", &config.client_id) + .append_pair("response_type", "code") + .append_pair("scope", &config.scopes.join(" ")) + .append_pair("code_challenge", challenge.as_str()) + .append_pair("code_challenge_method", "S256") + .append_pair("state", verifier.secret()); // ← Non-standard! + + if let Some(redirect_uri) = &config.redirect_uri { + url.query_pairs_mut() + .append_pair("redirect_uri", redirect_uri); + } + + if let Some(extra_params) = &config.extra_auth_params { + for (key, value) in extra_params { + url.query_pairs_mut().append_pair(key, value); + } + } + + Ok(AuthCodeParams { + auth_url: url.to_string(), + state: verifier.secret().to_string(), + code_verifier: Some(verifier.secret().to_string()), + }) + } + + async fn exchange_code( + &self, + config: &OAuthConfig, + code: &str, + verifier: Option<&str>, + ) -> anyhow::Result { + // Anthropic-specific token exchange + let (code, state) = if code.contains('#') { + let parts: Vec<&str> = code.split('#').collect(); + ( + parts.first().map(|s| s.to_string()).unwrap_or_default(), + parts.get(1).map(|s| s.to_string()), + ) + } else { + (code.to_string(), None) + }; + + let verifier = verifier + .ok_or_else(|| anyhow::anyhow!("PKCE verifier required for Anthropic OAuth"))?; + + let request_body = AnthropicTokenRequest { + code, + state: state.unwrap_or_else(|| verifier.to_string()), + grant_type: "authorization_code".to_string(), + client_id: config.client_id.to_string(), + redirect_uri: config.redirect_uri.clone(), + code_verifier: verifier.to_string(), + }; + + let client = self.build_http_client(config)?; + let response = client + .post(config.token_url.as_str()) + .header("Content-Type", "application/json") + .json(&request_body) + .send() + .await?; + + if !response.status().is_success() { + let status = response.status(); + let error_text = response.text().await.unwrap_or_default(); + anyhow::bail!("Token exchange failed with status {status}: {error_text}"); + } + + Ok(response.json().await?) + } + + /// Create HTTP client with provider-specific headers/behavior + fn build_http_client(&self, config: &OAuthConfig) -> anyhow::Result { + build_http_client(config.custom_headers.as_ref()) + } +} + +#[cfg(test)] +mod tests { + use forge_domain::OAuthConfig; + use url::Url; + + use super::*; + + fn test_oauth_config() -> OAuthConfig { + OAuthConfig { + client_id: "test_client".to_string().into(), + auth_url: Url::parse("https://example.com/auth").unwrap(), + token_url: Url::parse("https://example.com/token").unwrap(), + scopes: vec!["read".to_string(), "write".to_string()], + redirect_uri: Some("https://example.com/callback".to_string()), + use_pkce: true, + token_refresh_url: None, + extra_auth_params: None, + custom_headers: None, + } + } + + #[tokio::test] + async fn test_anthropic_provider_state_equals_verifier() { + let provider = AnthropicHttpProvider; + let config = test_oauth_config(); + + let result = provider.build_auth_url(&config).await.unwrap(); + + // Anthropic quirk: state should equal verifier + assert_eq!(result.state, result.code_verifier.unwrap()); + assert!(result.auth_url.contains("code_challenge_method=S256")); + } +} diff --git a/crates/forge_infra/src/auth/http/github.rs b/crates/forge_infra/src/auth/http/github.rs new file mode 100644 index 0000000000000000000000000000000000000000..3b3209febbff501fadaf65c03036ef5a684553bd --- /dev/null +++ b/crates/forge_infra/src/auth/http/github.rs @@ -0,0 +1,35 @@ +use forge_app::OAuthHttpProvider; +use forge_domain::{AuthCodeParams, OAuthConfig, OAuthTokenResponse}; + +use super::standard::StandardHttpProvider; +use crate::auth::util::build_http_client; + +/// GitHub Provider - HTTP 200 responses may contain errors +pub struct GithubHttpProvider; + +#[async_trait::async_trait] +impl OAuthHttpProvider for GithubHttpProvider { + async fn build_auth_url(&self, config: &OAuthConfig) -> anyhow::Result { + // Use standard flow - no quirks in auth URL + StandardHttpProvider.build_auth_url(config).await + } + + async fn exchange_code( + &self, + config: &OAuthConfig, + code: &str, + verifier: Option<&str>, + ) -> anyhow::Result { + // Use standard exchange - quirks handled in HTTP client via + // github_compliant_http_request + StandardHttpProvider + .exchange_code(config, code, verifier) + .await + } + + fn build_http_client(&self, config: &OAuthConfig) -> anyhow::Result { + // GitHub quirk: HTTP 200 responses may contain errors + // This is handled by the github_compliant_http_request function + build_http_client(config.custom_headers.as_ref()) + } +} diff --git a/crates/forge_infra/src/auth/http/mod.rs b/crates/forge_infra/src/auth/http/mod.rs new file mode 100644 index 0000000000000000000000000000000000000000..75fd44d7f0ff06e4a72454dafd965f3eafef238b --- /dev/null +++ b/crates/forge_infra/src/auth/http/mod.rs @@ -0,0 +1,7 @@ +mod anthropic; +mod github; +mod standard; + +pub(crate) use anthropic::AnthropicHttpProvider; +pub(crate) use github::GithubHttpProvider; +pub(crate) use standard::StandardHttpProvider; diff --git a/crates/forge_infra/src/auth/http/standard.rs b/crates/forge_infra/src/auth/http/standard.rs new file mode 100644 index 0000000000000000000000000000000000000000..162b307138e291cdc4bcfc3c863b0a4953083d66 --- /dev/null +++ b/crates/forge_infra/src/auth/http/standard.rs @@ -0,0 +1,139 @@ +use forge_app::OAuthHttpProvider; +use forge_domain::{AuthCodeParams, OAuthConfig, OAuthTokenResponse}; +use oauth2::{CsrfToken, PkceCodeChallenge, Scope}; +use serde::Serialize; + +use crate::auth::util::*; + +/// Standard RFC-compliant OAuth provider +pub struct StandardHttpProvider; + +#[derive(Debug, Serialize)] +struct StandardTokenRequest<'a> { + grant_type: &'static str, + code: &'a str, + client_id: &'a str, + #[serde(skip_serializing_if = "Option::is_none")] + redirect_uri: Option<&'a str>, + #[serde(skip_serializing_if = "Option::is_none")] + code_verifier: Option<&'a str>, +} + +#[async_trait::async_trait] +impl OAuthHttpProvider for StandardHttpProvider { + async fn build_auth_url(&self, config: &OAuthConfig) -> anyhow::Result { + // Use oauth2 library - standard flow + use oauth2::{AuthUrl, ClientId, TokenUrl}; + + let mut client = + oauth2::basic::BasicClient::new(ClientId::new(config.client_id.to_string())) + .set_auth_uri(AuthUrl::new(config.auth_url.to_string())?) + .set_token_uri(TokenUrl::new(config.token_url.to_string())?); + + if let Some(redirect_uri) = &config.redirect_uri { + client = client.set_redirect_uri(oauth2::RedirectUrl::new(redirect_uri.clone())?); + } + + let mut request = client.authorize_url(CsrfToken::new_random); + + for scope in &config.scopes { + request = request.add_scope(Scope::new(scope.clone())); + } + + if let Some(extra_params) = &config.extra_auth_params { + for (key, value) in extra_params { + request = request.add_extra_param(key, value); + } + } + + let (auth_url, csrf_state, pkce_verifier) = if config.use_pkce { + let (challenge, verifier) = PkceCodeChallenge::new_random_sha256(); + let (url, state) = request.set_pkce_challenge(challenge).url(); + (url, state, Some(verifier)) + } else { + let (url, state) = request.url(); + (url, state, None) + }; + + Ok(AuthCodeParams { + auth_url: auth_url.to_string(), + state: csrf_state.secret().to_string(), + code_verifier: pkce_verifier.map(|v| v.secret().to_string()), + }) + } + + async fn exchange_code( + &self, + config: &OAuthConfig, + code: &str, + verifier: Option<&str>, + ) -> anyhow::Result { + let http_client = self.build_http_client(config)?; + let request_body = StandardTokenRequest { + grant_type: "authorization_code", + code, + client_id: config.client_id.as_ref(), + redirect_uri: config.redirect_uri.as_deref(), + code_verifier: verifier, + }; + + let response = http_client + .post(config.token_url.as_str()) + .header("Content-Type", "application/x-www-form-urlencoded") + .header("Accept", "application/json") + .body(serde_urlencoded::to_string(&request_body)?) + .send() + .await?; + + let status = response.status(); + let body = response.text().await?; + + if !status.is_success() { + anyhow::bail!("OAuth token exchange failed ({status}): {body}"); + } + + // Parse the raw token payload so provider-specific fields like + // `id_token` are preserved instead of being dropped by generic helpers. + Ok(parse_token_response(&body)?) + } + + /// Create HTTP client with provider-specific headers/behavior + fn build_http_client(&self, config: &OAuthConfig) -> anyhow::Result { + build_http_client(config.custom_headers.as_ref()) + } +} + +#[cfg(test)] +mod tests { + use forge_domain::OAuthConfig; + use url::Url; + + use super::*; + + fn test_oauth_config() -> OAuthConfig { + OAuthConfig { + client_id: "test_client".to_string().into(), + auth_url: Url::parse("https://example.com/auth").unwrap(), + token_url: Url::parse("https://example.com/token").unwrap(), + scopes: vec!["read".to_string(), "write".to_string()], + redirect_uri: Some("https://example.com/callback".to_string()), + use_pkce: true, + token_refresh_url: None, + extra_auth_params: None, + custom_headers: None, + } + } + + #[tokio::test] + async fn test_standard_provider_build_auth_url() { + let provider = StandardHttpProvider; + let config = test_oauth_config(); + + let result = provider.build_auth_url(&config).await.unwrap(); + + assert!(result.auth_url.contains("client_id=test_client")); + assert!(result.auth_url.contains("response_type=code")); + assert!(result.code_verifier.is_some()); + assert_ne!(&result.state, result.code_verifier.as_ref().unwrap()); + } +} diff --git a/crates/forge_infra/src/auth/mcp_credentials.rs b/crates/forge_infra/src/auth/mcp_credentials.rs new file mode 100644 index 0000000000000000000000000000000000000000..e23fea5228a3aa562abe13d63dadf21738282cc6 --- /dev/null +++ b/crates/forge_infra/src/auth/mcp_credentials.rs @@ -0,0 +1,201 @@ +//! MCP OAuth Credential Storage +//! +//! Stores OAuth tokens separately from LLM provider credentials. +//! Credentials are bound to specific MCP server URLs. + +use std::collections::HashMap; + +use anyhow::Result; +use forge_domain::Environment; +use serde::{Deserialize, Serialize}; +use tokio::fs; + +/// MCP OAuth tokens for a single server. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct McpOAuthTokens { + pub access_token: String, + pub refresh_token: Option, + pub expires_at: Option, + pub scope: Option, +} + +/// Client registration info (for dynamic registration) +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct McpClientRegistration { + pub client_id: String, + pub client_secret: Option, + pub client_id_issued_at: Option, + pub client_secret_expires_at: Option, +} + +/// Complete credential entry for an MCP server +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct McpCredentialEntry { + /// Server URL (credential binding) + pub server_url: String, + /// OAuth tokens + pub tokens: McpOAuthTokens, + /// Client registration (if dynamically registered) + pub client_registration: Option, +} + +/// Credential store for all MCP servers +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct McpCredentialStore { + pub credentials: HashMap, +} + +impl McpCredentialStore { + /// Load the credential store from disk + /// + /// # Arguments + /// * `env` - The environment containing the base path for storage + /// + /// # Returns + /// * `Ok(McpCredentialStore)` - The loaded store, or empty if file doesn't + /// exist + /// * `Err(...)` - If there was an error reading or parsing the file + pub async fn load(env: &Environment) -> Result { + let path = Self::credential_path(env); + if !path.exists() { + return Ok(Self::default()); + } + let content = fs::read_to_string(&path).await?; + Ok(serde_json::from_str(&content)?) + } + + /// Save the credential store to disk + /// + /// Sets file permissions to 0o600 (user read/write only) on Unix systems. + /// + /// # Arguments + /// * `env` - The environment containing the base path for storage + pub async fn save(&self, env: &Environment) -> Result<()> { + let path = Self::credential_path(env); + if let Some(parent) = path.parent() { + fs::create_dir_all(parent).await?; + } + let content = serde_json::to_string_pretty(self)?; + fs::write(&path, content).await?; + + // Set permissions to 0o600 (user read/write only) + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + let metadata = fs::metadata(&path).await?; + let mut perms = metadata.permissions(); + perms.set_mode(0o600); + fs::set_permissions(&path, perms).await?; + } + Ok(()) + } + + /// Get the path to the credential file + /// + /// The file is stored at `/.mcp-credentials.json` + pub fn credential_path(env: &Environment) -> std::path::PathBuf { + env.base_path.join(".mcp-credentials.json") + } + + /// Get a credential entry by server URL + pub fn get(&self, server_url: &str) -> Option<&McpCredentialEntry> { + self.credentials.get(server_url) + } + + /// Set a credential entry + pub fn set(&mut self, entry: McpCredentialEntry) { + self.credentials.insert(entry.server_url.clone(), entry); + } + + /// Remove a credential entry by server URL + pub fn remove(&mut self, server_url: &str) { + self.credentials.remove(server_url); + } + + /// Check if credentials exist for a server URL + #[allow(dead_code)] + pub fn has_credentials(&self, server_url: &str) -> bool { + self.credentials.contains_key(server_url) + } +} + +#[cfg(test)] +mod tests { + use std::path::PathBuf; + + use tempfile::TempDir; + + use super::*; + + fn test_env(base_path: PathBuf) -> Environment { + Environment { + os: "test".to_string(), + cwd: PathBuf::from("/tmp"), + home: Some(PathBuf::from("/home/test")), + shell: "bash".to_string(), + base_path, + } + } + + #[tokio::test] + async fn test_credential_store_save_load() { + let tmp = TempDir::new().unwrap(); + let env = test_env(tmp.path().to_path_buf()); + let mut store = McpCredentialStore::default(); + + let entry = McpCredentialEntry { + server_url: "https://api.example.com/mcp".to_string(), + tokens: McpOAuthTokens { + access_token: "test-access-token".to_string(), + refresh_token: Some("test-refresh-token".to_string()), + expires_at: Some(1234567890), + scope: Some("read write".to_string()), + }, + client_registration: None, + }; + + store.set(entry); + store.save(&env).await.unwrap(); + + let loaded = McpCredentialStore::load(&env).await.unwrap(); + assert!(loaded.has_credentials("https://api.example.com/mcp")); + let entry = loaded.get("https://api.example.com/mcp").unwrap(); + assert_eq!(entry.tokens.access_token, "test-access-token"); + assert_eq!( + entry.tokens.refresh_token, + Some("test-refresh-token".to_string()) + ); + } + + #[tokio::test] + async fn test_credential_store_remove() { + let tmp = TempDir::new().unwrap(); + let env = test_env(tmp.path().to_path_buf()); + let mut store = McpCredentialStore::default(); + + store.set(McpCredentialEntry { + server_url: "https://api.example.com/mcp".to_string(), + tokens: McpOAuthTokens { + access_token: "test".to_string(), + refresh_token: None, + expires_at: None, + scope: None, + }, + client_registration: None, + }); + + store.remove("https://api.example.com/mcp"); + assert!(!store.has_credentials("https://api.example.com/mcp")); + + // Ensure the env reference is retained to keep the TempDir alive. + let _ = McpCredentialStore::credential_path(&env); + } + + #[tokio::test] + async fn test_credential_store_empty() { + let tmp = TempDir::new().unwrap(); + let env = test_env(tmp.path().to_path_buf()); + let store = McpCredentialStore::load(&env).await.unwrap(); + assert!(store.credentials.is_empty()); + } +} diff --git a/crates/forge_infra/src/auth/mcp_token_storage.rs b/crates/forge_infra/src/auth/mcp_token_storage.rs new file mode 100644 index 0000000000000000000000000000000000000000..50fde300d913b4d470fec5fb823ba0550420c344 --- /dev/null +++ b/crates/forge_infra/src/auth/mcp_token_storage.rs @@ -0,0 +1,275 @@ +//! Token Storage Adapter for rmcp OAuth +//! +//! Implements rmcp's `CredentialStore` trait to persist credentials +//! using our McpCredentialStore. + +use std::sync::Arc; + +use async_trait::async_trait; +use forge_domain::Environment; +use rmcp::transport::auth::{CredentialStore, StoredCredentials}; +use tokio::sync::Mutex; + +use crate::auth::mcp_credentials::{ + McpClientRegistration, McpCredentialEntry, McpCredentialStore, McpOAuthTokens, +}; + +/// Adapter that implements rmcp's CredentialStore trait +/// using our file-based McpCredentialStore +pub struct McpTokenStorage { + server_url: String, + env: Environment, + store: Arc>>, +} + +impl McpTokenStorage { + /// Create a new token storage adapter for a specific MCP server + /// + /// # Arguments + /// * `server_url` - The URL of the MCP server (used as credential key) + /// * `env` - The environment for file system paths + pub fn new(server_url: String, env: Environment) -> Self { + Self { server_url, env, store: Arc::new(Mutex::new(None)) } + } + + /// Get the credential store, loading it if necessary + async fn get_store(&self) -> anyhow::Result { + let mut guard = self.store.lock().await; + if guard.is_none() { + *guard = Some(McpCredentialStore::load(&self.env).await?); + } + Ok(guard.as_ref().unwrap().clone()) + } + + /// Update the cached store after modifications + async fn update_store(&self, store: McpCredentialStore) { + *self.store.lock().await = Some(store); + } + + /// Load stored credentials for this server + /// + /// Returns the stored credential entry if one exists, or None. + pub async fn load_credentials(&self) -> anyhow::Result> { + let store = self.get_store().await?; + Ok(store.get(&self.server_url).cloned()) + } + + /// Remove stored credentials for this server. + pub async fn remove_credentials(&self) -> anyhow::Result<()> { + let mut store = self.get_store().await?; + store.remove(&self.server_url); + store.save(&self.env).await?; + self.update_store(store).await; + Ok(()) + } + + /// Remove only tokens while keeping client registration. + /// Useful when tokens are expired/invalid but the client registration + /// (from dynamic registration) is still valid. + #[allow(dead_code)] + pub async fn remove_tokens_only(&self) -> anyhow::Result<()> { + let mut store = self.get_store().await?; + if let Some(entry) = store.get(&self.server_url).cloned() { + let updated = McpCredentialEntry { + server_url: entry.server_url, + tokens: McpOAuthTokens::default(), + client_registration: entry.client_registration, + }; + store.set(updated); + store.save(&self.env).await?; + self.update_store(store).await; + } + Ok(()) + } +} + +#[async_trait] +impl CredentialStore for McpTokenStorage { + /// Load credentials from storage + /// + /// Converts our file-based credentials to rmcp's StoredCredentials format. + /// Preserves refresh_token and expiry information for token refresh. + async fn load(&self) -> Result, rmcp::transport::auth::AuthError> { + let store = self + .get_store() + .await + .map_err(|e| rmcp::transport::auth::AuthError::InternalError(e.to_string()))?; + + if let Some(entry) = store.get(&self.server_url) { + use oauth2::basic::BasicTokenType; + use oauth2::{AccessToken, RefreshToken}; + use rmcp::transport::auth::{OAuthTokenResponse, VendorExtraTokenFields}; + + let access_token = AccessToken::new(entry.tokens.access_token.clone()); + let token_type = BasicTokenType::Bearer; + let extra_fields = VendorExtraTokenFields::default(); + + let mut token_response = + OAuthTokenResponse::new(access_token, token_type, extra_fields); + + // Set refresh token if available + if let Some(ref rt) = entry.tokens.refresh_token { + token_response.set_refresh_token(Some(RefreshToken::new(rt.clone()))); + } + + // Set expiry if available + if let Some(expires_at) = entry.tokens.expires_at { + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); + if expires_at > now { + token_response + .set_expires_in(Some(&std::time::Duration::from_secs(expires_at - now))); + } else { + // Token has expired - set zero duration so rmcp triggers refresh + token_response.set_expires_in(Some(&std::time::Duration::from_secs(0))); + } + } + + Ok(Some(StoredCredentials::new( + entry + .client_registration + .as_ref() + .map(|r| r.client_id.clone()) + .unwrap_or_default(), + Some(token_response), + vec![], + None, + ))) + } else { + Ok(None) + } + } + + /// Save credentials to storage + /// + /// Converts rmcp's StoredCredentials to our file-based format, + /// preserving all token metadata including refresh_token, expiry, + /// and client registration info. + async fn save( + &self, + credentials: StoredCredentials, + ) -> Result<(), rmcp::transport::auth::AuthError> { + use oauth2::TokenResponse; + + let mut store = self + .get_store() + .await + .map_err(|e| rmcp::transport::auth::AuthError::InternalError(e.to_string()))?; + + // Get existing entry to preserve client_secret and registration info + let existing = store.get(&self.server_url).cloned(); + + let tokens = if let Some(ref response) = credentials.token_response { + let access_token = response.access_token().secret().to_string(); + let refresh_token = response.refresh_token().map(|rt| rt.secret().to_string()); + let expires_at = response.expires_in().map(|d| { + (std::time::SystemTime::now() + d) + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs() + }); + let scope = response.scopes().map(|scopes| { + scopes + .iter() + .map(|s| s.to_string()) + .collect::>() + .join(" ") + }); + + McpOAuthTokens { + access_token, + refresh_token: refresh_token.or_else(|| { + existing + .as_ref() + .and_then(|e| e.tokens.refresh_token.clone()) + }), + expires_at, + scope, + } + } else { + McpOAuthTokens { + access_token: String::new(), + refresh_token: None, + expires_at: None, + scope: None, + } + }; + + // Preserve client_secret from dynamic registration + let client_registration = if credentials.client_id.is_empty() { + existing.and_then(|e| e.client_registration) + } else { + let existing_reg = existing + .as_ref() + .and_then(|e| e.client_registration.as_ref()); + Some(McpClientRegistration { + client_id: credentials.client_id, + client_secret: existing_reg.and_then(|r| r.client_secret.clone()), + client_id_issued_at: existing_reg.and_then(|r| r.client_id_issued_at), + client_secret_expires_at: existing_reg.and_then(|r| r.client_secret_expires_at), + }) + }; + + let entry = McpCredentialEntry { + server_url: self.server_url.clone(), + tokens, + client_registration, + }; + + store.set(entry); + store + .save(&self.env) + .await + .map_err(|e| rmcp::transport::auth::AuthError::InternalError(e.to_string()))?; + + self.update_store(store).await; + Ok(()) + } + + /// Clear credentials from storage + async fn clear(&self) -> Result<(), rmcp::transport::auth::AuthError> { + let mut store = self + .get_store() + .await + .map_err(|e| rmcp::transport::auth::AuthError::InternalError(e.to_string()))?; + + store.remove(&self.server_url); + store + .save(&self.env) + .await + .map_err(|e| rmcp::transport::auth::AuthError::InternalError(e.to_string()))?; + + self.update_store(store).await; + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use std::path::PathBuf; + + use super::*; + + fn test_env() -> Environment { + Environment { + os: "test".to_string(), + cwd: PathBuf::from("/tmp"), + home: Some(PathBuf::from("/home/test")), + shell: "bash".to_string(), + base_path: PathBuf::from("/tmp/test-forge"), + } + } + + #[tokio::test] + async fn test_token_storage_new() { + let env = test_env(); + let storage = McpTokenStorage::new("https://example.com/mcp".to_string(), env); + + // Should start with no credentials + let result = storage.load().await; + assert!(result.is_ok()); + assert!(result.unwrap().is_none()); + } +} diff --git a/crates/forge_infra/src/auth/mod.rs b/crates/forge_infra/src/auth/mod.rs new file mode 100644 index 0000000000000000000000000000000000000000..1edf7f732019e2ebbbd4af0c44c26da1f865bbda --- /dev/null +++ b/crates/forge_infra/src/auth/mod.rs @@ -0,0 +1,11 @@ +mod mcp_credentials; +mod mcp_token_storage; + +mod error; +mod http; +mod strategy; +mod util; + +pub(crate) use mcp_credentials::*; +pub(crate) use mcp_token_storage::*; +pub use strategy::*; diff --git a/crates/forge_infra/src/auth/strategy.rs b/crates/forge_infra/src/auth/strategy.rs new file mode 100644 index 0000000000000000000000000000000000000000..559a365f197fa0efd107d26eedefae0e7922d750 --- /dev/null +++ b/crates/forge_infra/src/auth/strategy.rs @@ -0,0 +1,1489 @@ +use std::time::Duration; + +use forge_app::{AuthStrategy, OAuthHttpProvider, StrategyFactory}; +use forge_domain::{ + ApiKey, ApiKeyRequest, AuthContextRequest, AuthContextResponse, AuthCredential, CodeRequest, + DeviceCodeRequest, OAuthConfig, OAuthTokenResponse, OAuthTokens, ProviderId, URLParam, + URLParamSpec, +}; +use google_cloud_auth::credentials::Builder; +use oauth2::basic::BasicClient; +use oauth2::{ClientId, DeviceAuthorizationUrl, Scope, TokenUrl}; +use reqwest::header::{HeaderMap, HeaderValue}; +use url::Url; + +use crate::auth::error::Error as AuthError; +use crate::auth::http::{AnthropicHttpProvider, GithubHttpProvider, StandardHttpProvider}; +use crate::auth::util::*; + +/// API Key Strategy - Simple static key authentication +pub struct ApiKeyStrategy { + provider_id: ProviderId, + required_params: Vec, +} + +impl ApiKeyStrategy { + pub fn new(provider_id: ProviderId, required_params: Vec) -> Self { + Self { provider_id, required_params } + } +} + +#[async_trait::async_trait] +impl AuthStrategy for ApiKeyStrategy { + async fn init(&self) -> anyhow::Result { + Ok(AuthContextRequest::ApiKey(ApiKeyRequest { + required_params: self.required_params.clone(), + existing_params: None, + api_key: None, + })) + } + + async fn complete( + &self, + context_response: AuthContextResponse, + ) -> anyhow::Result { + match context_response { + AuthContextResponse::ApiKey(ctx) => Ok(AuthCredential::new_api_key( + self.provider_id.clone(), + ctx.response.api_key, + ) + .url_params(ctx.response.url_params)), + _ => Err(AuthError::InvalidContext("Expected ApiKey context".to_string()).into()), + } + } + + async fn refresh(&self, credential: &AuthCredential) -> anyhow::Result { + // API keys don't expire - return as-is + Ok(credential.clone()) + } +} + +/// Extract the ChatGPT account ID from a JWT token's claims. +/// +/// Checks `chatgpt_account_id`, `https://api.openai.com/auth.chatgpt_account_id`, +/// and `organizations[0].id` in that order, matching the opencode +/// implementation. +fn extract_chatgpt_account_id(token: &str) -> Option { + let parts: Vec<&str> = token.split('.').collect(); + if parts.len() != 3 { + return None; + } + use base64::Engine; + let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(parts.get(1)?) + .ok()?; + let claims: serde_json::Value = serde_json::from_slice(&payload).ok()?; + + // Try chatgpt_account_id first + if let Some(id) = claims.get("chatgpt_account_id").and_then(|v| v.as_str()) { + return Some(id.to_string()); + } + // Try nested auth claim + if let Some(id) = claims + .get("https://api.openai.com/auth") + .and_then(|v| v.get("chatgpt_account_id")) + .and_then(|v| v.as_str()) + { + return Some(id.to_string()); + } + // Fall back to organizations[0].id + if let Some(id) = claims + .get("organizations") + .and_then(|v| v.as_array()) + .and_then(|orgs| orgs.first()) + .and_then(|org| org.get("id").and_then(|v| v.as_str())) + { + return Some(id.to_string()); + } + None +} + +/// Adds Codex-specific credential metadata derived from OAuth tokens. +/// +/// Tries to extract the account ID from the `id_token` first (which typically +/// contains the user identity claims in OpenID Connect flows), then falls back +/// to the `access_token` if needed. +fn enrich_codex_oauth_credential( + provider_id: &ProviderId, + credential: &mut AuthCredential, + id_token: Option<&str>, + access_token: &str, +) { + if *provider_id != ProviderId::CODEX { + return; + } + + // Try id_token first (preferred for user identity claims) + let account_id = id_token + .and_then(extract_chatgpt_account_id) + .or_else(|| extract_chatgpt_account_id(access_token)); + + if let Some(account_id) = account_id { + credential + .url_params + .insert("chatgpt_account_id".to_string().into(), account_id.into()); + } +} + +/// OAuth Code Strategy - Browser redirect flow +pub struct OAuthCodeStrategy { + provider_id: ProviderId, + config: OAuthConfig, + adapter: T, +} + +impl OAuthCodeStrategy { + pub fn new(adapter: T, provider_id: ProviderId, config: OAuthConfig) -> Self { + Self { config, provider_id, adapter } + } +} + +#[async_trait::async_trait] +impl AuthStrategy for OAuthCodeStrategy { + async fn init(&self) -> anyhow::Result { + let auth_params = self + .adapter + .build_auth_url(&self.config) + .await + .map_err(|e| AuthError::InitiationFailed(format!("Failed to build auth URL: {e}")))?; + + Ok(AuthContextRequest::Code(CodeRequest { + authorization_url: Url::parse(&auth_params.auth_url)?, + state: auth_params.state.into(), + pkce_verifier: auth_params.code_verifier.map(Into::into), + oauth_config: self.config.clone(), + })) + } + + async fn complete( + &self, + context_response: AuthContextResponse, + ) -> anyhow::Result { + match context_response { + AuthContextResponse::Code(ctx) => { + let token_response = self + .adapter + .exchange_code( + &ctx.request.oauth_config, + ctx.response.code.as_str(), + ctx.request.pkce_verifier.as_ref().map(|v| v.as_str()), + ) + .await + .map_err(|e| { + AuthError::CompletionFailed(format!( + "Failed to exchange authorization code: {e}" + )) + })?; + + let access_token = token_response.access_token.clone(); + let id_token = token_response.id_token.clone(); + let mut credential = build_oauth_credential( + self.provider_id.clone(), + token_response, + &ctx.request.oauth_config, + chrono::Duration::hours(1), // Code flow default + )?; + enrich_codex_oauth_credential( + &self.provider_id, + &mut credential, + id_token.as_deref(), + &access_token, + ); + Ok(credential) + } + _ => Err(AuthError::InvalidContext("Expected Code context".to_string()).into()), + } + } + + async fn refresh(&self, credential: &AuthCredential) -> anyhow::Result { + refresh_oauth_credential( + credential, + &self.config, + chrono::Duration::hours(1), + false, // No API key exchange + ) + .await + } +} + +/// OAuth Device Strategy - Device code flow +pub struct OAuthDeviceStrategy { + provider_id: ProviderId, + config: OAuthConfig, +} + +impl OAuthDeviceStrategy { + pub fn new(provider_id: ProviderId, config: OAuthConfig) -> Self { + Self { provider_id, config } + } +} + +#[async_trait::async_trait] +impl AuthStrategy for OAuthDeviceStrategy { + async fn init(&self) -> anyhow::Result { + // Build oauth2 client + let client = BasicClient::new(ClientId::new(self.config.client_id.to_string())) + .set_device_authorization_url( + DeviceAuthorizationUrl::new(self.config.auth_url.to_string()) + .map_err(|e| AuthError::InitiationFailed(format!("Invalid auth_url: {e}")))?, + ) + .set_token_uri( + TokenUrl::new(self.config.token_url.to_string()) + .map_err(|e| AuthError::InitiationFailed(format!("Invalid token_url: {e}")))?, + ); + + // Request device authorization + let mut request = client.exchange_device_code(); + for scope in &self.config.scopes { + request = request.add_scope(Scope::new(scope.clone())); + } + + // Build HTTP client with custom headers + let http_client = build_http_client(self.config.custom_headers.as_ref()).map_err(|e| { + AuthError::InitiationFailed(format!("Failed to build HTTP client: {e}")) + })?; + + let http_fn = |req| github_compliant_http_request(http_client.clone(), req); + + let device_auth_response: oauth2::StandardDeviceAuthorizationResponse = + request.request_async(&http_fn).await.map_err(|e| { + AuthError::InitiationFailed(format!("Device authorization request failed: {e}")) + })?; + + // Build the type-safe context + Ok(AuthContextRequest::DeviceCode(DeviceCodeRequest { + user_code: device_auth_response.user_code().secret().to_string().into(), + device_code: device_auth_response + .device_code() + .secret() + .to_string() + .into(), + verification_uri: Url::parse(&device_auth_response.verification_uri().to_string())?, + verification_uri_complete: device_auth_response + .verification_uri_complete() + .map(|u| Url::parse(&u.secret().to_string()).unwrap()), + expires_in: device_auth_response.expires_in().as_secs(), + interval: device_auth_response.interval().as_secs(), + oauth_config: self.config.clone(), + })) + } + + async fn complete( + &self, + context_response: AuthContextResponse, + ) -> anyhow::Result { + match context_response { + AuthContextResponse::DeviceCode(ctx) => { + let token_response = poll_for_tokens( + &ctx.request.device_code, + &self.config, + Duration::from_secs(600), + false, + ) + .await?; + + build_oauth_credential( + self.provider_id.clone(), + token_response, + &self.config, + chrono::Duration::days(365), // Device flow default + ) + } + _ => Err(AuthError::InvalidContext("Expected DeviceCode context".to_string()).into()), + } + } + + async fn refresh(&self, credential: &AuthCredential) -> anyhow::Result { + refresh_oauth_credential( + credential, + &self.config, + chrono::Duration::days(30), + false, // No API key exchange + ) + .await + } +} + +/// OAuth-with-API-Key Strategy - Hybrid flow (GitHub Copilot pattern) +pub struct OAuthWithApiKeyStrategy { + provider_id: ProviderId, + oauth_config: OAuthConfig, + api_key_exchange_url: Url, +} + +impl OAuthWithApiKeyStrategy { + pub fn new(provider_id: ProviderId, oauth_config: OAuthConfig) -> anyhow::Result { + let api_key_exchange_url = oauth_config + .token_refresh_url + .clone() + .ok_or_else(|| AuthError::InitiationFailed("Missing token_refresh_url".to_string()))?; + + Ok(Self { provider_id, oauth_config, api_key_exchange_url }) + } +} + +#[async_trait::async_trait] +impl AuthStrategy for OAuthWithApiKeyStrategy { + async fn init(&self) -> anyhow::Result { + // Same as OAuth Device init + let client = BasicClient::new(ClientId::new(self.oauth_config.client_id.to_string())) + .set_device_authorization_url( + DeviceAuthorizationUrl::new(self.oauth_config.auth_url.to_string()) + .map_err(|e| AuthError::InitiationFailed(format!("Invalid auth_url: {e}")))?, + ) + .set_token_uri( + TokenUrl::new(self.oauth_config.token_url.to_string()) + .map_err(|e| AuthError::InitiationFailed(format!("Invalid token_url: {e}")))?, + ); + + let mut request = client.exchange_device_code(); + for scope in &self.oauth_config.scopes { + request = request.add_scope(Scope::new(scope.clone())); + } + + let http_client = + build_http_client(self.oauth_config.custom_headers.as_ref()).map_err(|e| { + AuthError::InitiationFailed(format!("Failed to build HTTP client: {e}")) + })?; + + let http_fn = |req| github_compliant_http_request(http_client.clone(), req); + + let device_auth_response: oauth2::StandardDeviceAuthorizationResponse = + request.request_async(&http_fn).await.map_err(|e| { + AuthError::InitiationFailed(format!("Device authorization request failed: {e}")) + })?; + + Ok(AuthContextRequest::DeviceCode(DeviceCodeRequest { + user_code: device_auth_response.user_code().secret().to_string().into(), + device_code: device_auth_response + .device_code() + .secret() + .to_string() + .into(), + verification_uri: Url::parse(&device_auth_response.verification_uri().to_string())?, + verification_uri_complete: device_auth_response + .verification_uri_complete() + .map(|u| Url::parse(&u.secret().to_string()).unwrap()), + expires_in: device_auth_response.expires_in().as_secs(), + interval: device_auth_response.interval().as_secs(), + oauth_config: self.oauth_config.clone(), + })) + } + + async fn complete( + &self, + context_response: AuthContextResponse, + ) -> anyhow::Result { + match context_response { + AuthContextResponse::DeviceCode(ctx) => { + // Poll for OAuth tokens (GitHub-compatible) + let token_response = poll_for_tokens( + &ctx.request.device_code, + &self.oauth_config, + Duration::from_secs(600), + true, + ) + .await?; + + // Exchange for API key + let (api_key, expires_at) = exchange_oauth_for_api_key( + &token_response.access_token, + &self.api_key_exchange_url, + &self.oauth_config, + ) + .await?; + + let oauth_tokens = OAuthTokens::new( + token_response.access_token, + token_response.refresh_token, + expires_at, + ); + + Ok(AuthCredential::new_oauth_with_api_key( + self.provider_id.clone(), + oauth_tokens, + api_key, + self.oauth_config.clone(), + )) + } + _ => Err(AuthError::InvalidContext("Expected DeviceCode context".to_string()).into()), + } + } + + async fn refresh(&self, credential: &AuthCredential) -> anyhow::Result { + refresh_oauth_credential( + credential, + &self.oauth_config, + chrono::Duration::hours(1), // Unused for API key flow + true, // WITH API key exchange + ) + .await + } +} + +/// Google Application Default Credentials (ADC) Strategy +/// Uses Google Cloud SDK's ADC mechanism with automatic token refresh +pub struct GoogleAdcStrategy { + provider_id: ProviderId, + required_params: Vec, +} + +impl GoogleAdcStrategy { + pub fn new(provider_id: ProviderId, required_params: Vec) -> Self { + Self { provider_id, required_params } + } +} + +#[async_trait::async_trait] +impl AuthStrategy for GoogleAdcStrategy { + async fn init(&self) -> anyhow::Result { + // For Google ADC, we don't need any user interaction for the API key + // The credentials are automatically discovered from: + // 1. GOOGLE_APPLICATION_CREDENTIALS env var (service account) + // 2. gcloud ADC credentials (user credentials) + // 3. Metadata server (GCP environment) + // However, we still need to collect URL params like PROJECT_ID and LOCATION + Ok(AuthContextRequest::ApiKey(ApiKeyRequest { + required_params: self.required_params.clone(), + existing_params: None, + api_key: Some("google_adc_marker".to_string().into()), // Marker to indicate ADC usage + })) + } + + async fn complete( + &self, + context_response: AuthContextResponse, + ) -> anyhow::Result { + match context_response { + AuthContextResponse::ApiKey(ctx) => { + // Validate that gcloud auth is properly configured before completing + // authentication This ensures the user has run 'gcloud auth + // application-default login' + use google_cloud_auth::credentials::Builder; + const VERTEX_AI_SCOPES: &[&str] = + &["https://www.googleapis.com/auth/cloud-platform"]; + let credentials = Builder::default() + .with_scopes(VERTEX_AI_SCOPES.iter().map(|s| s.to_string())) + .build_access_token_credentials() + .map_err(|e| { + AuthError::CompletionFailed(format!( + "Google ADC not configured: {e}. Please run 'gcloud auth application-default login' to set up credentials." + )) + })?; + + // Try to fetch a token to verify authentication works + credentials + .access_token() + .await + .map_err(|e| { + AuthError::CompletionFailed(format!( + "Failed to obtain access token: {e}. Your ADC credentials may be expired — run 'gcloud auth application-default login' to re-authenticate." + )) + })?; + + // For Google ADC, we save a marker instead of the actual token + // The token will be refreshed on every use + // But we still need to save the url_params (PROJECT_ID, LOCATION) + Ok(AuthCredential::new_google_adc( + self.provider_id.clone(), + ApiKey::from("google_adc_marker".to_string()), /* Marker that will trigger + * refresh */ + ) + .url_params(ctx.response.url_params)) + } + _ => Err(AuthError::InvalidContext("Expected ApiKey context".to_string()).into()), + } + } + + async fn refresh(&self, credential: &AuthCredential) -> anyhow::Result { + // Google ADC handles token refresh automatically + // We just need to get a fresh token using the Builder API + // Vertex AI requires the cloud-platform scope + const VERTEX_AI_SCOPES: &[&str] = &["https://www.googleapis.com/auth/cloud-platform"]; + let credentials = Builder::default() + .with_scopes(VERTEX_AI_SCOPES.iter().map(|s| s.to_string())) + .build_access_token_credentials() + .map_err(|e| { + AuthError::RefreshFailed(format!( + "Failed to create Google credentials builder: {e}" + )) + })?; + + let access_token = credentials.access_token().await.map_err(|e| { + AuthError::RefreshFailed(format!( + "Failed to refresh Google access token: {e}. Your ADC credentials may be expired — run 'gcloud auth application-default login' to re-authenticate." + )) + })?; + + Ok(AuthCredential::new_google_adc( + self.provider_id.clone(), + ApiKey::from(access_token.token), + ) + .url_params(credential.url_params.clone())) + } +} + +/// AWS Profile Strategy - Uses AWS SDK credential chain with a named profile +/// Supports SSO, IAM, and other credential types configured in ~/.aws/config +pub struct AwsProfileStrategy { + provider_id: ProviderId, + required_params: Vec, +} + +const AWS_PROFILE_PARAM: &str = "AWS_PROFILE"; + +impl AwsProfileStrategy { + pub fn new(provider_id: ProviderId, mut required_params: Vec) -> Self { + let profile_param = URLParamSpec::new(URLParam::from(AWS_PROFILE_PARAM.to_string())); + if !required_params.iter().any(|p| p.name == profile_param.name) { + required_params.push(profile_param); + } + Self { provider_id, required_params } + } +} + +#[async_trait::async_trait] +impl AuthStrategy for AwsProfileStrategy { + async fn init(&self) -> anyhow::Result { + Ok(AuthContextRequest::ApiKey(ApiKeyRequest { + required_params: self.required_params.clone(), + existing_params: None, + api_key: Some("aws_profile_marker".to_string().into()), + })) + } + + async fn complete( + &self, + context_response: AuthContextResponse, + ) -> anyhow::Result { + match context_response { + AuthContextResponse::ApiKey(ctx) => { + let profile = ctx + .response + .url_params + .get(&URLParam::from(AWS_PROFILE_PARAM.to_string())) + .map(|v| v.to_string()) + .ok_or_else(|| { + AuthError::CompletionFailed("AWS_PROFILE is required".to_string()) + })?; + + // Validate the profile works by attempting to load credentials + let aws_config = aws_config::from_env().profile_name(&profile).load().await; + + let credentials_provider = + aws_config.credentials_provider().ok_or_else(|| { + AuthError::CompletionFailed(format!( + "No credentials found for profile '{}'. Ensure the profile exists in ~/.aws/config and you've run 'aws sso login --profile {}'", + profile, profile + )) + })?; + + // Try to resolve credentials to verify they work + use aws_credential_types::provider::ProvideCredentials; + credentials_provider + .provide_credentials() + .await + .map_err(|e| { + AuthError::CompletionFailed(format!( + "Failed to resolve credentials for profile '{}': {}. Try running 'aws sso login --profile {}'", + profile, e, profile + )) + })?; + + Ok( + AuthCredential::new_aws_profile( + self.provider_id.clone(), + ApiKey::from(profile), + ) + .url_params(ctx.response.url_params), + ) + } + _ => Err(AuthError::InvalidContext("Expected ApiKey context".to_string()).into()), + } + } + + async fn refresh(&self, credential: &AuthCredential) -> anyhow::Result { + // AWS SDK handles SSO token refresh internally + Ok(credential.clone()) + } +} + +/// OpenAI Codex Device Strategy - Custom device auth for ChatGPT Pro/Plus +/// +/// Implements the OpenAI-specific device authorization flow used by Codex: +/// 1. Request device code from `/api/accounts/deviceauth/usercode` +/// 2. User enters code at `https://auth.openai.com/codex/device` +/// 3. Poll `/api/accounts/deviceauth/token` for authorization code + verifier +/// 4. Exchange authorization code for OAuth tokens via standard token endpoint +pub struct CodexDeviceStrategy { + provider_id: ProviderId, + config: OAuthConfig, +} + +impl CodexDeviceStrategy { + pub fn new(provider_id: ProviderId, config: OAuthConfig) -> Self { + Self { provider_id, config } + } +} + +/// Response from the OpenAI device auth usercode endpoint +#[derive(Debug, serde::Deserialize)] +struct CodexDeviceAuthResponse { + device_auth_id: String, + user_code: String, + interval: String, +} + +/// Response from the OpenAI device auth token polling endpoint +#[derive(Debug, serde::Deserialize)] +struct CodexDeviceTokenResponse { + authorization_code: String, + code_verifier: String, +} + +#[async_trait::async_trait] +impl AuthStrategy for CodexDeviceStrategy { + async fn init(&self) -> anyhow::Result { + let http_client = build_http_client(self.config.custom_headers.as_ref()).map_err(|e| { + AuthError::InitiationFailed(format!("Failed to build HTTP client: {e}")) + })?; + + // Step 1: Request device authorization from OpenAI's custom endpoint + let response = http_client + .post(self.config.auth_url.as_str()) + .header("Content-Type", "application/json") + .json(&serde_json::json!({ + "client_id": self.config.client_id.as_str() + })) + .send() + .await + .map_err(|e| { + AuthError::InitiationFailed(format!("Device authorization request failed: {e}")) + })?; + + if !response.status().is_success() { + return Err(AuthError::InitiationFailed(format!( + "Device authorization failed with status: {}", + response.status() + )) + .into()); + } + + let device_data: CodexDeviceAuthResponse = response.json().await.map_err(|e| { + AuthError::InitiationFailed(format!("Failed to parse device auth response: {e}")) + })?; + + let interval: u64 = device_data.interval.parse().unwrap_or(5).max(1); + + // Build the device code request using existing domain types + // We encode the device_auth_id in the device_code field + Ok(AuthContextRequest::DeviceCode(DeviceCodeRequest { + user_code: device_data.user_code.clone().into(), + device_code: device_data.device_auth_id.into(), + verification_uri: Url::parse("https://auth.openai.com/codex/device")?, + verification_uri_complete: None, + expires_in: 300, // 5 minute timeout + interval, + oauth_config: self.config.clone(), + })) + } + + async fn complete( + &self, + context_response: AuthContextResponse, + ) -> anyhow::Result { + match context_response { + AuthContextResponse::DeviceCode(ctx) => { + // Poll for authorization code using the custom OpenAI endpoint + let token_response = codex_poll_for_tokens(&ctx.request, &self.config).await?; + + let access_token = token_response.access_token.clone(); + let id_token = token_response.id_token.clone(); + let mut credential = build_oauth_credential( + self.provider_id.clone(), + token_response, + &self.config, + chrono::Duration::hours(1), + )?; + + // Store account_id in url_params so it's persisted and available + // for chat request headers. + enrich_codex_oauth_credential( + &self.provider_id, + &mut credential, + id_token.as_deref(), + &access_token, + ); + + Ok(credential) + } + _ => Err(AuthError::InvalidContext("Expected DeviceCode context".to_string()).into()), + } + } + + async fn refresh(&self, credential: &AuthCredential) -> anyhow::Result { + refresh_oauth_credential(credential, &self.config, chrono::Duration::hours(1), false).await + } +} + +/// Refresh OAuth credential - handles all OAuth flows +async fn refresh_oauth_credential( + credential: &AuthCredential, + config: &OAuthConfig, + expiry_duration: chrono::Duration, + with_api_key_exchange: bool, +) -> anyhow::Result { + // Extract tokens (works for OAuth and OAuthWithApiKey) + let tokens = extract_oauth_tokens(credential)?; + + // Determine which OAuth access token to use + let (oauth_access_token, oauth_refresh_token) = + if let Some(refresh_token) = &tokens.refresh_token { + // If we have a refresh token, refresh the OAuth access token first + tracing::debug!("Refreshing OAuth access token using refresh token"); + let token_response = refresh_access_token(config, refresh_token.as_str()).await?; + ( + token_response.access_token.clone(), + token_response.refresh_token, + ) + } else { + // No refresh token - use the existing long-lived OAuth access token + // This is typical for GitHub Copilot where the OAuth token is long-lived + tracing::debug!("No refresh token available, using existing OAuth access token"); + ( + tokens.access_token.to_string(), + tokens.refresh_token.clone().map(|t| t.to_string()), + ) + }; + + // Exchange for API key if needed (GitHub Copilot pattern) + let (api_key, expires_at) = if with_api_key_exchange { + let url = config.token_refresh_url.as_ref().ok_or_else(|| { + AuthError::RefreshFailed("Missing token_refresh_url for API key exchange".to_string()) + })?; + let (key, expiry) = exchange_oauth_for_api_key(&oauth_access_token, url, config).await?; + (Some(key), expiry) + } else { + let expiry = calculate_token_expiry(None, expiry_duration); + (None, expiry) + }; + + // Build new tokens with refreshed OAuth access token + let new_tokens = OAuthTokens::new(oauth_access_token, oauth_refresh_token, expires_at); + + // Build appropriate credential type while preserving URL params + let refreshed = if let Some(key) = api_key { + AuthCredential::new_oauth_with_api_key( + credential.id.clone(), + new_tokens, + key, + config.clone(), + ) + } else { + AuthCredential::new_oauth(credential.id.clone(), new_tokens, config.clone()) + }; + + Ok(refreshed.url_params(credential.url_params.clone())) +} + +/// Poll for OAuth tokens during device flow +async fn poll_for_tokens( + device_code: &forge_domain::DeviceCode, + config: &OAuthConfig, + timeout: Duration, + github_compatible: bool, +) -> anyhow::Result { + let http_client = build_http_client(config.custom_headers.as_ref()) + .map_err(|e| AuthError::PollFailed(format!("Failed to build HTTP client: {e}")))?; + + let start_time = tokio::time::Instant::now(); + let interval = Duration::from_secs(5); + + loop { + // Check timeout + if start_time.elapsed() >= timeout { + return Err(AuthError::Timeout(timeout).into()); + } + + // Sleep before polling (GitHub pattern only) + if github_compatible { + tokio::time::sleep(interval).await; + } + + // Build token request + let params = vec![ + ( + "grant_type".to_string(), + "urn:ietf:params:oauth:grant-type:device_code".to_string(), + ), + ("device_code".to_string(), device_code.to_string()), + ("client_id".to_string(), config.client_id.to_string()), + ]; + + let body = serde_urlencoded::to_string(¶ms) + .map_err(|e| AuthError::PollFailed(format!("Failed to encode request: {e}")))?; + + // Make HTTP request with headers + let mut headers = HeaderMap::new(); + headers.insert( + "Content-Type", + HeaderValue::from_static("application/x-www-form-urlencoded"), + ); + headers.insert("Accept", HeaderValue::from_static("application/json")); + + inject_custom_headers(&mut headers, &config.custom_headers); + + let response = http_client + .post(config.token_url.as_str()) + .headers(headers) + .body(body) + .send() + .await + .map_err(|e| AuthError::PollFailed(format!("HTTP request failed: {e}")))?; + + let status = response.status(); + let body_text = response + .text() + .await + .map_err(|e| AuthError::PollFailed(format!("Failed to read response: {e}")))?; + + // GitHub-compatible: HTTP 200 can contain either success or error + if github_compatible && status.is_success() { + let token_response: serde_json::Value = serde_json::from_str(&body_text) + .unwrap_or_else(|_| serde_json::json!({"error": "parse_error"})); + + // Check for error field first + if let Some(error) = token_response.get("error").and_then(|v| v.as_str()) { + if handle_oauth_error(error).is_ok() { + // Retryable error - continue polling + continue; + } + // Terminal error - propagate + return Err(handle_oauth_error(error).unwrap_err().into()); + } + + // No error field - parse as success + return Ok(parse_token_response(&body_text)?); + } + + // Standard OAuth: HTTP success means tokens + if !github_compatible && status.is_success() { + return Ok(parse_token_response(&body_text)?); + } + + // Handle error responses (non-200 status for standard OAuth) + let error_response: serde_json::Value = serde_json::from_str(&body_text) + .unwrap_or_else(|_| serde_json::json!({"error": "unknown_error"})); + + if let Some(error) = error_response.get("error").and_then(|v| v.as_str()) { + if handle_oauth_error(error).is_ok() { + // Retryable error - sleep and continue + tokio::time::sleep(if error == "slow_down" { + interval * 2 + } else { + interval + }) + .await; + continue; + } + // Terminal error - propagate + return Err(handle_oauth_error(error).unwrap_err().into()); + } + + // Unknown error + return Err(AuthError::PollFailed(format!("HTTP {status}: {body_text}")).into()); + } +} + +/// Poll for Codex tokens using OpenAI's custom device auth endpoints. +/// +/// This differs from standard OAuth2 device code flow: +/// 1. Polls `/api/accounts/deviceauth/token` with `device_auth_id` + +/// `user_code` +/// 2. Receives `authorization_code` + `code_verifier` (not tokens directly) +/// 3. Exchanges the authorization code for OAuth tokens via standard token +/// endpoint +async fn codex_poll_for_tokens( + request: &DeviceCodeRequest, + config: &OAuthConfig, +) -> anyhow::Result { + let http_client = build_http_client(config.custom_headers.as_ref()) + .map_err(|e| AuthError::PollFailed(format!("Failed to build HTTP client: {e}")))?; + + let timeout = Duration::from_secs(request.expires_in); + let interval = Duration::from_secs(request.interval.max(1)); + // Add a safety margin to polling interval to avoid rate limiting + let poll_interval = interval + Duration::from_secs(3); + + let start_time = tokio::time::Instant::now(); + + // The auth_url in config points to the usercode endpoint; derive the token + // polling endpoint from the same base + let poll_url = config.auth_url.as_str().replace("/usercode", "/token"); + + loop { + if start_time.elapsed() >= timeout { + return Err(AuthError::Timeout(timeout).into()); + } + + // Poll request and response handling would be here + + tokio::time::sleep(poll_interval).await; + + let response = http_client + .post(&poll_url) + .header("Content-Type", "application/json") + .json(&serde_json::json!({ + "device_auth_id": request.device_code.as_str(), + "user_code": request.user_code.as_str(), + })) + .send() + .await + .map_err(|e| AuthError::PollFailed(format!("HTTP request failed: {e}")))?; + + let status = response.status(); + + if status.is_success() { + // Parse the custom response containing authorization_code + code_verifier + let device_token: CodexDeviceTokenResponse = response.json().await.map_err(|e| { + AuthError::PollFailed(format!("Failed to parse device token response: {e}")) + })?; + + // Exchange the authorization code for OAuth tokens via standard + // endpoint. Use a clean HTTP client without custom headers since the + // standard OAuth token endpoint rejects unknown headers. + let clean_client = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .build() + .map_err(|e| AuthError::PollFailed(format!("Failed to build HTTP client: {e}")))?; + + let token_response = clean_client + .post(config.token_url.as_str()) + .header("Content-Type", "application/x-www-form-urlencoded") + .body( + serde_urlencoded::to_string([ + ("grant_type", "authorization_code"), + ("code", &device_token.authorization_code), + ( + "redirect_uri", + "https://auth.openai.com/deviceauth/callback", + ), + ("client_id", config.client_id.as_ref()), + ("code_verifier", &device_token.code_verifier), + ]) + .map_err(|e| { + AuthError::PollFailed(format!("Failed to encode token request: {e}")) + })?, + ) + .send() + .await + .map_err(|e| { + AuthError::PollFailed(format!("Token exchange request failed: {e}")) + })?; + + if !token_response.status().is_success() { + let token_exchange_status = token_response.status(); + let error_text = token_response.text().await.unwrap_or_default(); + return Err(AuthError::PollFailed(format!( + "Token exchange failed ({}): {}", + token_exchange_status, error_text + )) + .into()); + } + + return Ok(parse_token_response( + &token_response.text().await.map_err(|e| { + AuthError::PollFailed(format!("Failed to read token response: {e}")) + })?, + )?); + } + + // 403/404 means authorization pending (user hasn't entered code yet) + if status.as_u16() == 403 || status.as_u16() == 404 { + continue; + } + + // Any other error is terminal + let body_text = response.text().await.unwrap_or_default(); + return Err(AuthError::PollFailed(format!("HTTP {status}: {body_text}")).into()); + } +} + +/// Exchange OAuth token for API key (GitHub Copilot pattern) +async fn exchange_oauth_for_api_key( + oauth_token: &str, + api_key_exchange_url: &Url, + config: &OAuthConfig, +) -> anyhow::Result<(ApiKey, chrono::DateTime)> { + // Build request headers + let mut headers = reqwest::header::HeaderMap::new(); + headers.insert( + reqwest::header::AUTHORIZATION, + reqwest::header::HeaderValue::from_str(&format!("Bearer {oauth_token}")).map_err(|e| { + AuthError::CompletionFailed(format!("Invalid authorization header: {e}")) + })?, + ); + + // Add custom headers from config + inject_custom_headers(&mut headers, &config.custom_headers); + + let response = build_http_client(config.custom_headers.as_ref()) + .map_err(|e| AuthError::CompletionFailed(format!("Failed to build HTTP client: {e}")))? + .get(api_key_exchange_url.as_str()) + .headers(headers) + .send() + .await + .map_err(|e| { + AuthError::CompletionFailed(format!("API key exchange request failed: {e}")) + })?; + + let status = response.status(); + if !status.is_success() { + if status.as_u16() == 403 { + return Err(AuthError::CompletionFailed( + "Access denied. Ensure you have an active subscription.".to_string(), + ) + .into()); + } + return Err(AuthError::CompletionFailed(format!( + "API key fetch failed ({}): {}", + status, + response.text().await.unwrap_or_default() + )) + .into()); + } + + let OAuthTokenResponse { access_token, expires_at, .. } = + response.json().await.map_err(|e| { + AuthError::CompletionFailed(format!("Failed to parse API key response: {e}")) + })?; + + Ok(( + access_token.into(), + chrono::DateTime::from_timestamp(expires_at.unwrap_or(0), 0) + .unwrap_or_else(chrono::Utc::now), + )) +} + +/// Enum wrapper for all strategy implementations +/// Eliminates heap allocation and dynamic dispatch +pub enum AnyAuthStrategy { + ApiKey(ApiKeyStrategy), + OAuthCodeStandard(OAuthCodeStrategy), + OAuthCodeAnthropic(OAuthCodeStrategy), + OAuthCodeGithub(OAuthCodeStrategy), + OAuthDevice(OAuthDeviceStrategy), + OAuthWithApiKey(OAuthWithApiKeyStrategy), + GoogleAdc(GoogleAdcStrategy), + AwsProfile(AwsProfileStrategy), + CodexDevice(CodexDeviceStrategy), +} + +#[async_trait::async_trait] +impl AuthStrategy for AnyAuthStrategy { + async fn init(&self) -> anyhow::Result { + match self { + Self::ApiKey(s) => s.init().await, + Self::OAuthCodeStandard(s) => s.init().await, + Self::OAuthCodeAnthropic(s) => s.init().await, + Self::OAuthCodeGithub(s) => s.init().await, + Self::OAuthDevice(s) => s.init().await, + Self::OAuthWithApiKey(s) => s.init().await, + Self::GoogleAdc(s) => s.init().await, + Self::AwsProfile(s) => s.init().await, + Self::CodexDevice(s) => s.init().await, + } + } + + async fn complete( + &self, + context_response: AuthContextResponse, + ) -> anyhow::Result { + match self { + Self::ApiKey(s) => s.complete(context_response).await, + Self::OAuthCodeStandard(s) => s.complete(context_response).await, + Self::OAuthCodeAnthropic(s) => s.complete(context_response).await, + Self::OAuthCodeGithub(s) => s.complete(context_response).await, + Self::OAuthDevice(s) => s.complete(context_response).await, + Self::OAuthWithApiKey(s) => s.complete(context_response).await, + Self::GoogleAdc(s) => s.complete(context_response).await, + Self::AwsProfile(s) => s.complete(context_response).await, + Self::CodexDevice(s) => s.complete(context_response).await, + } + } + + async fn refresh(&self, credential: &AuthCredential) -> anyhow::Result { + match self { + Self::ApiKey(s) => s.refresh(credential).await, + Self::OAuthCodeStandard(s) => s.refresh(credential).await, + Self::OAuthCodeAnthropic(s) => s.refresh(credential).await, + Self::OAuthCodeGithub(s) => s.refresh(credential).await, + Self::OAuthDevice(s) => s.refresh(credential).await, + Self::OAuthWithApiKey(s) => s.refresh(credential).await, + Self::GoogleAdc(s) => s.refresh(credential).await, + Self::AwsProfile(s) => s.refresh(credential).await, + Self::CodexDevice(s) => s.refresh(credential).await, + } + } +} + +/// Factory for creating authentication strategies +pub struct ForgeAuthStrategyFactory; + +impl Default for ForgeAuthStrategyFactory { + fn default() -> Self { + Self + } +} + +impl ForgeAuthStrategyFactory { + pub fn new(_environment: forge_domain::Environment) -> Self { + Self + } +} + +impl StrategyFactory for ForgeAuthStrategyFactory { + type Strategy = AnyAuthStrategy; + + fn create_auth_strategy( + &self, + provider_id: ProviderId, + auth_method: forge_domain::AuthMethod, + required_params: Vec, + ) -> anyhow::Result { + match auth_method { + forge_domain::AuthMethod::ApiKey => Ok(AnyAuthStrategy::ApiKey(ApiKeyStrategy::new( + provider_id, + required_params, + ))), + forge_domain::AuthMethod::OAuthCode(config) => { + if provider_id == ProviderId::CLAUDE_CODE { + return Ok(AnyAuthStrategy::OAuthCodeAnthropic(OAuthCodeStrategy::new( + AnthropicHttpProvider, + provider_id, + config, + ))); + } + + if provider_id == ProviderId::GITHUB_COPILOT { + return Ok(AnyAuthStrategy::OAuthCodeGithub(OAuthCodeStrategy::new( + GithubHttpProvider, + provider_id, + config, + ))); + } + + Ok(AnyAuthStrategy::OAuthCodeStandard(OAuthCodeStrategy::new( + StandardHttpProvider, + provider_id, + config, + ))) + } + forge_domain::AuthMethod::OAuthDevice(config) => { + // Check if this is OAuth-with-API-Key flow (GitHub Copilot pattern) + if config.token_refresh_url.is_some() { + Ok(AnyAuthStrategy::OAuthWithApiKey( + OAuthWithApiKeyStrategy::new(provider_id, config)?, + )) + } else { + Ok(AnyAuthStrategy::OAuthDevice(OAuthDeviceStrategy::new( + provider_id, + config, + ))) + } + } + forge_domain::AuthMethod::GoogleAdc => Ok(AnyAuthStrategy::GoogleAdc( + GoogleAdcStrategy::new(provider_id, required_params), + )), + forge_domain::AuthMethod::AwsProfile => Ok(AnyAuthStrategy::AwsProfile( + AwsProfileStrategy::new(provider_id, required_params), + )), + forge_domain::AuthMethod::CodexDevice(config) => Ok(AnyAuthStrategy::CodexDevice( + CodexDeviceStrategy::new(provider_id, config), + )), + } + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use forge_domain::URLParam; + use pretty_assertions::assert_eq; + + use super::*; + + #[test] + fn test_create_auth_strategy_api_key() { + let factory = ForgeAuthStrategyFactory; + let strategy = factory.create_auth_strategy( + ProviderId::OPENAI, + forge_domain::AuthMethod::ApiKey, + vec![], + ); + assert!(strategy.is_ok()); + } + + #[test] + fn test_create_auth_strategy_oauth_code() { + let config = OAuthConfig { + client_id: "test".to_string().into(), + auth_url: Url::parse("https://example.com/auth").unwrap(), + token_url: Url::parse("https://example.com/token").unwrap(), + scopes: vec![], + redirect_uri: None, + use_pkce: false, + token_refresh_url: None, + extra_auth_params: None, + custom_headers: None, + }; + + let factory = ForgeAuthStrategyFactory; + let strategy = factory.create_auth_strategy( + ProviderId::OPENAI, + forge_domain::AuthMethod::OAuthCode(config), + vec![], + ); + assert!(strategy.is_ok()); + } + + #[test] + fn test_create_auth_strategy_oauth_device() { + let config = OAuthConfig { + client_id: "test".to_string().into(), + auth_url: Url::parse("https://example.com/auth").unwrap(), + token_url: Url::parse("https://example.com/token").unwrap(), + scopes: vec![], + redirect_uri: None, + use_pkce: false, + token_refresh_url: None, + extra_auth_params: None, + custom_headers: None, + }; + + let factory = ForgeAuthStrategyFactory; + let strategy = factory.create_auth_strategy( + ProviderId::OPENAI, + forge_domain::AuthMethod::OAuthDevice(config), + vec![], + ); + assert!(strategy.is_ok()); + } + + #[test] + fn test_create_auth_strategy_oauth_with_api_key() { + let config = OAuthConfig { + client_id: "test".to_string().into(), + auth_url: Url::parse("https://example.com/auth").unwrap(), + token_url: Url::parse("https://example.com/token").unwrap(), + scopes: vec![], + redirect_uri: None, + use_pkce: false, + token_refresh_url: Some(Url::parse("https://example.com/api_key").unwrap()), + extra_auth_params: None, + custom_headers: None, + }; + + let factory = ForgeAuthStrategyFactory; + let strategy = factory.create_auth_strategy( + ProviderId::GITHUB_COPILOT, + forge_domain::AuthMethod::OAuthDevice(config), + vec![], + ); + assert!(strategy.is_ok()); + } + + #[test] + fn test_create_auth_strategy_codex_device() { + let config = OAuthConfig { + client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string().into(), + auth_url: Url::parse("https://auth.openai.com/api/accounts/deviceauth/usercode") + .unwrap(), + token_url: Url::parse("https://auth.openai.com/oauth/token").unwrap(), + scopes: vec![], + redirect_uri: None, + use_pkce: false, + token_refresh_url: None, + extra_auth_params: None, + custom_headers: None, + }; + + let factory = ForgeAuthStrategyFactory; + let actual = factory.create_auth_strategy( + ProviderId::CODEX, + forge_domain::AuthMethod::CodexDevice(config), + vec![], + ); + assert!(actual.is_ok()); + assert!(matches!(actual.unwrap(), AnyAuthStrategy::CodexDevice(_))); + } + + /// Helper to build a JWT token with the given claims payload. + fn build_jwt(claims: &serde_json::Value) -> String { + use base64::Engine; + let header = base64::engine::general_purpose::URL_SAFE_NO_PAD + .encode(r#"{"alg":"RS256","typ":"JWT"}"#); + let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD + .encode(serde_json::to_vec(claims).unwrap()); + format!("{header}.{payload}.fake_signature") + } + + #[test] + fn test_extract_chatgpt_account_id_from_direct_claim() { + let fixture = build_jwt(&serde_json::json!({ + "chatgpt_account_id": "acct_123" + })); + let actual = extract_chatgpt_account_id(&fixture); + let expected = Some("acct_123".to_string()); + assert_eq!(actual, expected); + } + + #[test] + fn test_extract_chatgpt_account_id_from_nested_claim() { + let fixture = build_jwt(&serde_json::json!({ + "https://api.openai.com/auth": { + "chatgpt_account_id": "acct_nested_456" + } + })); + let actual = extract_chatgpt_account_id(&fixture); + let expected = Some("acct_nested_456".to_string()); + assert_eq!(actual, expected); + } + + #[test] + fn test_extract_chatgpt_account_id_from_organizations() { + let fixture = build_jwt(&serde_json::json!({ + "organizations": [ + {"id": "org_789", "name": "My Org"} + ] + })); + let actual = extract_chatgpt_account_id(&fixture); + let expected = Some("org_789".to_string()); + assert_eq!(actual, expected); + } + + #[test] + fn test_extract_chatgpt_account_id_prefers_direct_claim() { + let fixture = build_jwt(&serde_json::json!({ + "chatgpt_account_id": "direct", + "https://api.openai.com/auth": { + "chatgpt_account_id": "nested" + }, + "organizations": [{"id": "org"}] + })); + let actual = extract_chatgpt_account_id(&fixture); + let expected = Some("direct".to_string()); + assert_eq!(actual, expected); + } + + #[test] + fn test_extract_chatgpt_account_id_returns_none_for_empty_claims() { + let fixture = build_jwt(&serde_json::json!({})); + let actual = extract_chatgpt_account_id(&fixture); + assert_eq!(actual, None); + } + + #[test] + fn test_extract_chatgpt_account_id_returns_none_for_invalid_jwt() { + let actual = extract_chatgpt_account_id("not-a-jwt"); + assert_eq!(actual, None); + } + + #[test] + fn test_extract_chatgpt_account_id_returns_none_for_invalid_base64() { + let actual = extract_chatgpt_account_id("header.!!!invalid-base64!!!.signature"); + assert_eq!(actual, None); + } + + #[test] + fn test_extract_chatgpt_account_id_returns_none_for_empty_organizations() { + let fixture = build_jwt(&serde_json::json!({ + "organizations": [] + })); + let actual = extract_chatgpt_account_id(&fixture); + assert_eq!(actual, None); + } + + #[test] + fn test_enrich_codex_oauth_credential_uses_id_token_claims() { + let fixture_id_token = build_jwt(&serde_json::json!({ + "chatgpt_account_id": "acct_from_id_token" + })); + let fixture_access_token = "not-a-jwt"; + let mut actual = AuthCredential::new_oauth( + ProviderId::CODEX, + OAuthTokens::new( + fixture_access_token, + None::, + chrono::Utc::now() + chrono::Duration::hours(1), + ), + OAuthConfig { + client_id: "test".to_string().into(), + auth_url: Url::parse("https://example.com/auth").unwrap(), + token_url: Url::parse("https://example.com/token").unwrap(), + scopes: vec![], + redirect_uri: Some("http://localhost:1455/auth/callback".to_string()), + use_pkce: true, + token_refresh_url: None, + extra_auth_params: None, + custom_headers: None, + }, + ); + + enrich_codex_oauth_credential( + &ProviderId::CODEX, + &mut actual, + Some(&fixture_id_token), + fixture_access_token, + ); + + let actual = actual + .url_params + .get(&URLParam::from("chatgpt_account_id".to_string())); + let expected = Some(&forge_domain::URLParamValue::from( + "acct_from_id_token".to_string(), + )); + + assert_eq!(actual, expected); + } + + #[tokio::test] + async fn test_refresh_oauth_credential_preserves_url_params() { + let fixture_config = OAuthConfig { + client_id: "test".to_string().into(), + auth_url: Url::parse("https://example.com/auth").unwrap(), + token_url: Url::parse("https://example.com/token").unwrap(), + scopes: vec![], + redirect_uri: None, + use_pkce: false, + token_refresh_url: None, + extra_auth_params: None, + custom_headers: None, + }; + let fixture_tokens = OAuthTokens::new( + "access_token", + None::, + chrono::Utc::now() + chrono::Duration::minutes(30), + ); + let fixture_url_params = HashMap::from([( + URLParam::from("chatgpt_account_id".to_string()), + "acct_123".to_string().into(), + )]); + let fixture_credential = + AuthCredential::new_oauth(ProviderId::CODEX, fixture_tokens, fixture_config.clone()) + .url_params(fixture_url_params.clone()); + + let actual = refresh_oauth_credential( + &fixture_credential, + &fixture_config, + chrono::Duration::hours(1), + false, + ) + .await + .unwrap(); + + let expected = fixture_url_params; + assert_eq!(actual.url_params, expected); + } +} diff --git a/crates/forge_infra/src/auth/util.rs b/crates/forge_infra/src/auth/util.rs new file mode 100644 index 0000000000000000000000000000000000000000..a3890fc6b0f6d7aac0dc68cba25cc796a9d9f68d --- /dev/null +++ b/crates/forge_infra/src/auth/util.rs @@ -0,0 +1,281 @@ +use std::collections::HashMap; + +use chrono::Utc; +use forge_domain::{ + AuthCredential, AuthDetails, OAuthConfig, OAuthTokenResponse, OAuthTokens, ProviderId, +}; +use oauth2::basic::BasicClient; +use oauth2::{ClientId, RefreshToken, TokenUrl}; + +use crate::auth::error::Error; + +/// Calculate token expiry with fallback duration +pub(crate) fn calculate_token_expiry( + expires_in: Option, + fallback: chrono::Duration, +) -> chrono::DateTime { + if let Some(seconds) = expires_in { + Utc::now() + chrono::Duration::seconds(seconds as i64) + } else { + Utc::now() + fallback + } +} + +/// Convert oauth2 TokenResponse into domain OAuthTokenResponse +pub(crate) fn into_domain(token: T) -> OAuthTokenResponse { + OAuthTokenResponse { + access_token: token.access_token().secret().to_string(), + refresh_token: token.refresh_token().map(|t| t.secret().to_string()), + expires_in: token.expires_in().map(|d| d.as_secs()), + expires_at: None, + token_type: "Bearer".to_string(), + scope: token.scopes().map(|scopes| { + scopes + .iter() + .map(|s| s.to_string()) + .collect::>() + .join(" ") + }), + id_token: None, // oauth2 crate doesn't provide id_token directly + } +} + +/// Build HTTP client with custom headers +pub(crate) fn build_http_client( + custom_headers: Option<&HashMap>, +) -> anyhow::Result { + let mut builder = reqwest::Client::builder() + // Disable redirects to prevent SSRF vulnerabilities + .redirect(reqwest::redirect::Policy::none()); + + if let Some(headers) = custom_headers { + let mut header_map = reqwest::header::HeaderMap::new(); + + for (key, value) in headers { + let header_name = reqwest::header::HeaderName::try_from(key.as_str()) + .map_err(|e| anyhow::anyhow!("Invalid header name '{key}': {e}"))?; + let header_value = value + .parse() + .map_err(|e| anyhow::anyhow!("Invalid header value for '{key}': {e}"))?; + header_map.insert(header_name, header_value); + } + + builder = builder.default_headers(header_map); + } + + Ok(builder.build()?) +} + +/// Build OAuth credential with consistent expiry handling +pub(crate) fn build_oauth_credential( + provider_id: ProviderId, + token_response: OAuthTokenResponse, + config: &OAuthConfig, + default_expiry: chrono::Duration, +) -> anyhow::Result { + let expires_at = calculate_token_expiry(token_response.expires_in, default_expiry); + let oauth_tokens = OAuthTokens::new( + token_response.access_token, + token_response.refresh_token, + expires_at, + ); + Ok(AuthCredential::new_oauth( + provider_id, + oauth_tokens, + config.clone(), + )) +} + +/// Extract OAuth tokens from any credential type +pub(crate) fn extract_oauth_tokens(credential: &AuthCredential) -> anyhow::Result<&OAuthTokens> { + match &credential.auth_details { + AuthDetails::OAuth { tokens, .. } => Ok(tokens), + AuthDetails::OAuthWithApiKey { tokens, .. } => Ok(tokens), + _ => Err( + Error::RefreshFailed("Invalid credential type for token extraction".to_string()).into(), + ), + } +} + +/// Refresh OAuth access token using refresh token +pub(crate) async fn refresh_access_token( + config: &OAuthConfig, + refresh_token: &str, +) -> anyhow::Result { + // Build minimal oauth2 client (just need token endpoint) + let client = BasicClient::new(ClientId::new(config.client_id.to_string())) + .set_token_uri(TokenUrl::new(config.token_url.to_string())?); + + // Build HTTP client with custom headers + let http_client = build_http_client(config.custom_headers.as_ref())?; + + let refresh_token = RefreshToken::new(refresh_token.to_string()); + + // Use GitHub-compliant HTTP function to handle non-RFC responses + let http_fn = |req| github_compliant_http_request(http_client.clone(), req); + + let token_result = client + .exchange_refresh_token(&refresh_token) + .request_async(&http_fn) + .await?; + + Ok(into_domain(token_result)) +} + +/// GitHub-compliant HTTP request handler that fixes status codes for error +/// responses +pub(crate) async fn github_compliant_http_request( + client: reqwest::Client, + request: http::Request>, +) -> Result>, reqwest::Error> { + // Execute the request + let mut req_builder = client + .request(request.method().clone(), request.uri().to_string()) + .body(request.body().clone()); + + for (name, value) in request.headers() { + req_builder = req_builder.header(name.as_str(), value.as_bytes()); + } + + let response = req_builder.send().await?; + + // Get status and body + let status_code = response.status(); + let headers = response.headers().clone(); + let body = response.bytes().await?; + + // GitHub-specific fix: If status is 200 but body contains "error" field, + // change status to 400 so oauth2 crate recognizes it as an error response + let fixed_status = if status_code.is_success() { + if let Ok(json) = serde_json::from_slice::(&body) { + if json.get("error").is_some() { + // This is actually an error response masquerading as success + http::StatusCode::BAD_REQUEST + } else { + status_code + } + } else { + status_code + } + } else { + status_code + }; + + // Build http::Response with corrected status + let mut response_builder = http::Response::builder().status(fixed_status); + + // Add headers + for (name, value) in headers.iter() { + response_builder = response_builder.header(name, value); + } + + Ok(response_builder + .body(body.to_vec()) + .expect("Failed to build HTTP response")) +} + +/// Inject custom headers into a header map +pub(crate) fn inject_custom_headers( + headers: &mut reqwest::header::HeaderMap, + custom_headers: &Option>, +) { + use reqwest::header::{HeaderName, HeaderValue}; + + if let Some(custom_headers) = custom_headers { + for (key, value) in custom_headers { + if let (Ok(name), Ok(val)) = (HeaderName::try_from(key), HeaderValue::from_str(value)) { + headers.insert(name, val); + } + } + } +} + +/// Parse OAuth error responses during polling +pub(crate) fn handle_oauth_error(error_code: &str) -> Result<(), Error> { + match error_code { + "authorization_pending" | "slow_down" => Ok(()), + "expired_token" => Err(Error::Expired), + "access_denied" => Err(Error::Denied), + _ => Err(Error::PollFailed(format!("OAuth error: {error_code}"))), + } +} + +/// Parse token response from JSON. +pub(crate) fn parse_token_response(body: &str) -> Result { + let token_response: OAuthTokenResponse = serde_json::from_str(body) + .map_err(|e| Error::PollFailed(format!("Failed to parse token response: {e}")))?; + + if token_response.access_token.trim().is_empty() { + return Err(Error::PollFailed( + "Missing access_token in response".to_string(), + )); + } + + Ok(token_response) +} + +#[cfg(test)] +mod tests { + use chrono::Duration; + + use super::*; + + #[test] + fn test_calculate_token_expiry_with_expires_in() { + let before = Utc::now(); + let expires_at = calculate_token_expiry(Some(3600), Duration::hours(1)); + let after = Utc::now() + Duration::hours(1); + + assert!(expires_at >= before); + assert!(expires_at <= after); + } + + #[test] + fn test_calculate_token_expiry_with_fallback() { + let before = Utc::now(); + let expires_at = calculate_token_expiry(None, Duration::days(365)); + let after = Utc::now() + Duration::days(365); + + assert!(expires_at >= before); + assert!(expires_at <= after); + } + + #[test] + fn test_parse_token_response_preserves_id_token() { + let fixture = r#"{ + "access_token": "test_token", + "refresh_token": "refresh_token", + "expires_in": 3600, + "id_token": "test_id_token" + }"#; + + let actual = parse_token_response(fixture).unwrap(); + + assert_eq!(actual.access_token, "test_token"); + assert_eq!(actual.refresh_token, Some("refresh_token".to_string())); + assert_eq!(actual.expires_in, Some(3600)); + assert_eq!(actual.id_token, Some("test_id_token".to_string())); + } + + #[test] + fn test_handle_oauth_error_retryable() { + assert!(handle_oauth_error("authorization_pending").is_ok()); + assert!(handle_oauth_error("slow_down").is_ok()); + } + + #[test] + fn test_handle_oauth_error_terminal() { + assert!(matches!( + handle_oauth_error("expired_token"), + Err(Error::Expired) + )); + assert!(matches!( + handle_oauth_error("access_denied"), + Err(Error::Denied) + )); + assert!(matches!( + handle_oauth_error("unknown_error"), + Err(Error::PollFailed(_)) + )); + } +} diff --git a/crates/forge_infra/src/executor.rs b/crates/forge_infra/src/executor.rs new file mode 100644 index 0000000000000000000000000000000000000000..af8718738d18697331a047c7c9f83d58c21402ca --- /dev/null +++ b/crates/forge_infra/src/executor.rs @@ -0,0 +1,538 @@ +use std::io::{self, Write}; +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +use bstr::ByteSlice; +use forge_app::CommandInfra; +use forge_domain::{CommandOutput, ConsoleWriter as OutputPrinterTrait, Environment}; +use tokio::io::AsyncReadExt; +use tokio::process::Command; +use tokio::sync::Mutex; + +use crate::console::StdConsoleWriter; + +/// Service for executing shell commands +#[derive(Clone, Debug)] +pub struct ForgeCommandExecutorService { + env: Environment, + output_printer: Arc, + + // Mutex to ensure that only one command is executed at a time + ready: Arc>, +} + +impl ForgeCommandExecutorService { + pub fn new(env: Environment, output_printer: Arc) -> Self { + Self { env, output_printer, ready: Arc::new(Mutex::new(())) } + } + + fn prepare_command( + &self, + command_str: &str, + working_dir: &Path, + env_vars: Option>, + ) -> Command { + // Create a basic command + let is_windows = cfg!(target_os = "windows"); + let shell = self.env.shell.as_str(); + let mut command = Command::new(shell); + + // Core color settings for general commands + command + .env("CLICOLOR_FORCE", "1") + .env("FORCE_COLOR", "true") + .env_remove("NO_COLOR"); + + // Language/program specific color settings + command + .env("SBT_OPTS", "-Dsbt.color=always") + .env("JAVA_OPTS", "-Dsbt.color=always"); + + // enabled Git colors + command.env("GIT_CONFIG_PARAMETERS", "'color.ui=always'"); + + // Other common tools + command.env("GREP_OPTIONS", "--color=always"); // GNU grep + + let parameter = if is_windows { "/C" } else { "-c" }; + command.arg(parameter); + + #[cfg(windows)] + command.raw_arg(command_str); + #[cfg(unix)] + command.arg(command_str); + + tracing::info!(command = command_str, "Executing command"); + + command.kill_on_drop(true); + + // Set the working directory + command.current_dir(working_dir); + + // Configure the command for output + command + .stdin(std::process::Stdio::inherit()) + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::piped()); + + // Set requested environment variables + if let Some(env_vars) = env_vars { + for env_var in env_vars { + if let Ok(value) = std::env::var(&env_var) { + command.env(&env_var, value); + tracing::debug!(env_var = %env_var, "Set environment variable from system"); + } else { + tracing::warn!(env_var = %env_var, "Environment variable not found in system"); + } + } + } + + command + } + + /// Internal method to execute commands with streaming to console + async fn execute_command_internal( + &self, + command: String, + working_dir: &Path, + silent: bool, + env_vars: Option>, + ) -> anyhow::Result { + let ready = self.ready.lock().await; + + let mut prepared_command = self.prepare_command(&command, working_dir, env_vars); + + // Spawn the command + let mut child = prepared_command.spawn()?; + + let mut stdout_pipe = child.stdout.take(); + let mut stderr_pipe = child.stderr.take(); + + // Stream the output of the command to stdout and stderr concurrently + let (status, stdout_buffer, stderr_buffer) = if silent { + tokio::try_join!( + child.wait(), + stream(&mut stdout_pipe, io::sink()), + stream(&mut stderr_pipe, io::sink()) + )? + } else { + let stdout_writer = OutputPrinterWriter::stdout(self.output_printer.clone()); + let stderr_writer = OutputPrinterWriter::stderr(self.output_printer.clone()); + let result = tokio::try_join!( + child.wait(), + stream(&mut stdout_pipe, stdout_writer), + stream(&mut stderr_pipe, stderr_writer) + )?; + + // If the command's stdout did not end with a newline, the terminal + // cursor is left mid-line. Write a newline so that subsequent output + // (e.g. the LLM response) starts on a fresh line. + if result.1.last() != Some(&b'\n') && !result.1.is_empty() { + let _ = self.output_printer.write(b"\n"); + let _ = self.output_printer.flush(); + } + + result + }; + + // Drop happens after `try_join` due to + drop(stdout_pipe); + drop(stderr_pipe); + drop(ready); + + Ok(CommandOutput { + stdout: stdout_buffer.to_str_lossy().into_owned(), + stderr: stderr_buffer.to_str_lossy().into_owned(), + exit_code: status.code(), + command, + }) + } +} + +/// Writer that delegates to OutputPrinter for synchronized writes. +struct OutputPrinterWriter { + printer: Arc, + is_stdout: bool, +} + +impl OutputPrinterWriter { + fn stdout(printer: Arc) -> Self { + Self { printer, is_stdout: true } + } + + fn stderr(printer: Arc) -> Self { + Self { printer, is_stdout: false } + } +} + +impl Write for OutputPrinterWriter { + fn write(&mut self, buf: &[u8]) -> io::Result { + if self.is_stdout { + self.printer.write(buf) + } else { + self.printer.write_err(buf) + } + } + + fn flush(&mut self) -> io::Result<()> { + if self.is_stdout { + self.printer.flush() + } else { + self.printer.flush_err() + } + } +} + +/// reads the output from A and writes it to W +async fn stream( + io: &mut Option, + mut writer: W, +) -> io::Result> { + let mut output = Vec::new(); + if let Some(io) = io.as_mut() { + let mut buff = [0; 1024]; + // Carry incomplete trailing UTF-8 codepoint bytes across reads — Windows + // console stdio rejects even one byte of a split codepoint. + let mut pending = Vec::::new(); + loop { + let n = io.read(&mut buff).await?; + if n == 0 { + break; + } + let chunk = buff.get(..n).unwrap_or(&[]); + output.extend_from_slice(chunk); + + let mut working = std::mem::take(&mut pending); + working.extend_from_slice(chunk); + pending = write_lossy_utf8(&mut writer, &working)?; + // note: flush is necessary else we get the cursor could not be found error. + writer.flush()?; + } + // Flush dangling bytes from a stream that ended mid-codepoint. + if !pending.is_empty() { + writer.write_all(pending.to_str_lossy().as_bytes())?; + writer.flush()?; + } + } + Ok(output) +} + +/// Writes `buf` as valid UTF-8 (invalid bytes → `U+FFFD`) and returns any +/// incomplete trailing codepoint bytes for the caller to carry into the next +/// chunk. +fn write_lossy_utf8(writer: &mut W, buf: &[u8]) -> io::Result> { + let mut chunks = ByteSlice::utf8_chunks(buf).peekable(); + + while let Some(chunk) = chunks.next() { + writer.write_all(chunk.valid().as_bytes())?; + + if !chunk.invalid().is_empty() { + if chunk.incomplete() && chunks.peek().is_none() { + return Ok(chunk.invalid().to_vec()); + } + writer.write_all("\u{FFFD}".as_bytes())?; + } + } + + Ok(Vec::new()) +} + +/// The implementation for CommandExecutorService +#[async_trait::async_trait] +impl CommandInfra for ForgeCommandExecutorService { + async fn execute_command( + &self, + command: String, + working_dir: PathBuf, + silent: bool, + env_vars: Option>, + ) -> anyhow::Result { + self.execute_command_internal(command, &working_dir, silent, env_vars) + .await + } + + async fn execute_command_raw( + &self, + command: &str, + working_dir: PathBuf, + env_vars: Option>, + ) -> anyhow::Result { + let mut prepared_command = self.prepare_command(command, &working_dir, env_vars); + + // overwrite the stdin, stdout and stderr to inherit + prepared_command + .stdin(std::process::Stdio::inherit()) + .stdout(std::process::Stdio::inherit()) + .stderr(std::process::Stdio::inherit()); + + Ok(prepared_command.spawn()?.wait().await?) + } +} + +#[cfg(test)] +mod tests { + + use pretty_assertions::assert_eq; + + use super::*; + + fn test_env() -> Environment { + use fake::{Fake, Faker}; + let fixture: Environment = Faker.fake(); + fixture.shell( + if cfg!(target_os = "windows") { + "cmd" + } else { + "bash" + } + .to_string(), + ) + } + + fn test_printer() -> Arc { + Arc::new(StdConsoleWriter::default()) + } + + #[tokio::test] + async fn test_command_executor() { + let fixture = ForgeCommandExecutorService::new(test_env(), test_printer()); + let cmd = "echo 'hello world'"; + let dir = "."; + + let actual = fixture + .execute_command(cmd.to_string(), PathBuf::new().join(dir), false, None) + .await + .unwrap(); + + let mut expected = CommandOutput { + stdout: "hello world\n".to_string(), + stderr: "".to_string(), + command: "echo \"hello world\"".into(), + exit_code: Some(0), + }; + + if cfg!(target_os = "windows") { + expected.stdout = format!("'{}'", expected.stdout); + } + + assert_eq!(actual.stdout.trim(), expected.stdout.trim()); + assert_eq!(actual.stderr, expected.stderr); + assert_eq!(actual.success(), expected.success()); + } + #[tokio::test] + async fn test_command_executor_with_env_vars_success() { + // Set up test environment variables + unsafe { + std::env::set_var("TEST_ENV_VAR", "test_value"); + std::env::set_var("ANOTHER_TEST_VAR", "another_value"); + } + + let fixture = ForgeCommandExecutorService::new(test_env(), test_printer()); + let cmd = if cfg!(target_os = "windows") { + "echo %TEST_ENV_VAR%" + } else { + "echo $TEST_ENV_VAR" + }; + + let actual = fixture + .execute_command( + cmd.to_string(), + PathBuf::new().join("."), + false, + Some(vec!["TEST_ENV_VAR".to_string()]), + ) + .await + .unwrap(); + + assert!(actual.success()); + assert!(actual.stdout.contains("test_value")); + + // Clean up + unsafe { + std::env::remove_var("TEST_ENV_VAR"); + std::env::remove_var("ANOTHER_TEST_VAR"); + } + } + + #[tokio::test] + async fn test_command_executor_with_missing_env_vars() { + unsafe { + std::env::remove_var("MISSING_ENV_VAR"); + } + + let fixture = ForgeCommandExecutorService::new(test_env(), test_printer()); + let cmd = if cfg!(target_os = "windows") { + "echo %MISSING_ENV_VAR%" + } else { + "echo ${MISSING_ENV_VAR:-default_value}" + }; + + let actual = fixture + .execute_command( + cmd.to_string(), + PathBuf::new().join("."), + false, + Some(vec!["MISSING_ENV_VAR".to_string()]), + ) + .await + .unwrap(); + + // Should still succeed even with missing env vars + assert!(actual.success()); + } + + #[tokio::test] + async fn test_command_executor_with_empty_env_list() { + let fixture = ForgeCommandExecutorService::new(test_env(), test_printer()); + let cmd = "echo 'no env vars'"; + + let actual = fixture + .execute_command( + cmd.to_string(), + PathBuf::new().join("."), + false, + Some(vec![]), + ) + .await + .unwrap(); + + assert!(actual.success()); + assert!(actual.stdout.contains("no env vars")); + } + + #[tokio::test] + async fn test_command_executor_with_multiple_env_vars() { + unsafe { + std::env::set_var("FIRST_VAR", "first"); + std::env::set_var("SECOND_VAR", "second"); + } + + let fixture = ForgeCommandExecutorService::new(test_env(), test_printer()); + let cmd = if cfg!(target_os = "windows") { + "echo %FIRST_VAR% %SECOND_VAR%" + } else { + "echo $FIRST_VAR $SECOND_VAR" + }; + + let actual = fixture + .execute_command( + cmd.to_string(), + PathBuf::new().join("."), + false, + Some(vec!["FIRST_VAR".to_string(), "SECOND_VAR".to_string()]), + ) + .await + .unwrap(); + + assert!(actual.success()); + assert!(actual.stdout.contains("first")); + assert!(actual.stdout.contains("second")); + + // Clean up + unsafe { + std::env::remove_var("FIRST_VAR"); + std::env::remove_var("SECOND_VAR"); + } + } + + #[tokio::test] + async fn test_command_executor_silent() { + let fixture = ForgeCommandExecutorService::new(test_env(), test_printer()); + let cmd = "echo 'silent test'"; + let dir = "."; + + let actual = fixture + .execute_command(cmd.to_string(), PathBuf::new().join(dir), true, None) + .await + .unwrap(); + + let mut expected = CommandOutput { + stdout: "silent test\n".to_string(), + stderr: "".to_string(), + command: "echo \"silent test\"".into(), + exit_code: Some(0), + }; + + if cfg!(target_os = "windows") { + expected.stdout = format!("'{}'", expected.stdout); + } + + // The output should still be captured in the CommandOutput + assert_eq!(actual.stdout.trim(), expected.stdout.trim()); + assert_eq!(actual.stderr, expected.stderr); + assert_eq!(actual.success(), expected.success()); + } + + mod write_lossy_utf8 { + use pretty_assertions::assert_eq; + + use super::super::write_lossy_utf8; + + fn run(buf: &[u8]) -> (Vec, Vec) { + let mut out = Vec::::new(); + let pending = write_lossy_utf8(&mut out, buf).unwrap(); + (out, pending) + } + + #[test] + fn valid_ascii_passes_through() { + let (out, pending) = run(b"hello"); + assert_eq!(out, b"hello"); + assert!(pending.is_empty()); + } + + #[test] + fn valid_multibyte_passes_through() { + // "héllo ✓" — mixed 2-byte and 3-byte codepoints. + let input = "héllo ✓".as_bytes(); + let (out, pending) = run(input); + assert_eq!(out, input); + assert!(pending.is_empty()); + } + + #[test] + fn incomplete_trailing_codepoint_is_buffered() { + // "é" is 0xC3 0xA9 — leading byte alone must be held back. + let (out, pending) = run(&[b'a', 0xC3]); + assert_eq!(out, b"a"); + assert_eq!(pending, vec![0xC3]); + } + + #[test] + fn multibyte_split_across_two_chunks_emits_once_whole() { + let mut out = Vec::::new(); + let pending = write_lossy_utf8(&mut out, &[b'a', 0xC3]).unwrap(); + assert_eq!(pending, vec![0xC3]); + assert_eq!(out, b"a"); + + let mut working = pending; + working.push(0xA9); + let pending = write_lossy_utf8(&mut out, &working).unwrap(); + assert!(pending.is_empty()); + assert_eq!(out, "aé".as_bytes()); + } + + #[test] + fn invalid_byte_in_middle_becomes_replacement() { + let (out, pending) = run(&[b'a', 0xFF, b'b']); + assert_eq!(out, "a\u{FFFD}b".as_bytes()); + assert!(pending.is_empty()); + } + + #[test] + fn lone_continuation_byte_becomes_replacement() { + let (out, pending) = run(&[b'a', 0x80, b'b']); + assert_eq!(out, "a\u{FFFD}b".as_bytes()); + assert!(pending.is_empty()); + } + + #[test] + fn windows_1252_smart_quote_becomes_replacement() { + // Regression: 0x91/0x92 land as bare continuation bytes and broke + // console stdio on Windows before this fix. + let (out, pending) = run(b"quote: \x91hi\x92"); + assert_eq!(out, "quote: \u{FFFD}hi\u{FFFD}".as_bytes()); + assert!(pending.is_empty()); + } + } +} diff --git a/crates/forge_infra/src/fs_create_dirs.rs b/crates/forge_infra/src/fs_create_dirs.rs new file mode 100644 index 0000000000000000000000000000000000000000..40b63e418191058a74389887a9c1ac90bffd826f --- /dev/null +++ b/crates/forge_infra/src/fs_create_dirs.rs @@ -0,0 +1,13 @@ +use std::path::Path; + +use forge_app::FileDirectoryInfra; + +#[derive(Default)] +pub struct ForgeCreateDirsService; + +#[async_trait::async_trait] +impl FileDirectoryInfra for ForgeCreateDirsService { + async fn create_dirs(&self, path: &Path) -> anyhow::Result<()> { + Ok(forge_fs::ForgeFS::create_dir_all(path).await?) + } +} diff --git a/crates/forge_infra/src/fs_meta.rs b/crates/forge_infra/src/fs_meta.rs new file mode 100644 index 0000000000000000000000000000000000000000..16d3505367bc1c8f7600384ebb7cba06da366e44 --- /dev/null +++ b/crates/forge_infra/src/fs_meta.rs @@ -0,0 +1,24 @@ +use std::path::Path; + +use anyhow::Result; +use forge_app::FileInfoInfra; + +pub struct ForgeFileMetaService; +#[async_trait::async_trait] +impl FileInfoInfra for ForgeFileMetaService { + async fn is_file(&self, path: &Path) -> Result { + Ok(forge_fs::ForgeFS::is_file(path)) + } + + async fn is_binary(&self, path: &Path) -> Result { + forge_fs::ForgeFS::is_binary_file(path).await + } + + async fn exists(&self, path: &Path) -> Result { + Ok(forge_fs::ForgeFS::exists(path)) + } + + async fn file_size(&self, path: &Path) -> Result { + forge_fs::ForgeFS::file_size(path).await + } +} diff --git a/crates/forge_infra/src/fs_read.rs b/crates/forge_infra/src/fs_read.rs new file mode 100644 index 0000000000000000000000000000000000000000..408b8da732811333b9b424d3f7c83668160c84d2 --- /dev/null +++ b/crates/forge_infra/src/fs_read.rs @@ -0,0 +1,137 @@ +use std::path::{Path, PathBuf}; + +use anyhow::Result; +use forge_app::FileReaderInfra; +use futures::{StreamExt, stream}; + +#[derive(Clone)] +pub struct ForgeFileReadService; + +impl Default for ForgeFileReadService { + fn default() -> Self { + Self + } +} + +impl ForgeFileReadService { + pub fn new() -> Self { + Self + } +} + +#[async_trait::async_trait] +impl FileReaderInfra for ForgeFileReadService { + async fn read_utf8(&self, path: &Path) -> Result { + forge_fs::ForgeFS::read_utf8(path).await + } + + fn read_batch_utf8( + &self, + batch_size: usize, + paths: Vec, + ) -> impl futures::Stream)> + Send { + let batches: Vec> = paths + .chunks(batch_size) + .map(|chunk| chunk.to_vec()) + .collect(); + + stream::iter(batches) + .then(move |batch| async move { + let futures = batch.into_iter().map(|path| async move { + let result = self.read_utf8(&path).await; + (path, result) + }); + + futures::future::join_all(futures).await + }) + .flat_map(stream::iter) + } + + async fn read(&self, path: &Path) -> Result> { + forge_fs::ForgeFS::read(path).await + } + + async fn range_read_utf8( + &self, + path: &Path, + start_line: u64, + end_line: u64, + ) -> Result<(String, forge_domain::FileInfo)> { + forge_fs::ForgeFS::read_range_utf8(path, start_line, end_line).await + } +} + +#[cfg(test)] +mod tests { + use std::io::Write; + + use futures::StreamExt; + use tempfile::NamedTempFile; + + use super::*; + + #[tokio::test] + async fn test_read_batch_utf8() { + let fixture = ForgeFileReadService::new(); + + // Create temporary test files + let mut file1 = NamedTempFile::new().unwrap(); + let mut file2 = NamedTempFile::new().unwrap(); + let mut file3 = NamedTempFile::new().unwrap(); + + writeln!(file1, "content1").unwrap(); + writeln!(file2, "content2").unwrap(); + writeln!(file3, "content3").unwrap(); + + let paths = vec![ + file1.path().to_path_buf(), + file2.path().to_path_buf(), + file3.path().to_path_buf(), + ]; + + // Read with batch size of 2 + let stream = fixture.read_batch_utf8(2, paths.clone()); + futures::pin_mut!(stream); + + let item1 = stream.next().await.unwrap(); + assert_eq!(item1.0, paths[0]); + assert_eq!(item1.1.as_deref().unwrap().trim(), "content1"); + + let item2 = stream.next().await.unwrap(); + assert_eq!(item2.0, paths[1]); + assert_eq!(item2.1.as_deref().unwrap().trim(), "content2"); + + let item3 = stream.next().await.unwrap(); + assert_eq!(item3.0, paths[2]); + assert_eq!(item3.1.as_deref().unwrap().trim(), "content3"); + + // No more items + assert!(stream.next().await.is_none()); + } + + #[tokio::test] + async fn test_read_batch_utf8_single_batch() { + let fixture = ForgeFileReadService::new(); + + let mut file1 = NamedTempFile::new().unwrap(); + let mut file2 = NamedTempFile::new().unwrap(); + + writeln!(file1, "test1").unwrap(); + writeln!(file2, "test2").unwrap(); + + let paths = vec![file1.path().to_path_buf(), file2.path().to_path_buf()]; + + // Read with batch size larger than number of files + let stream = fixture.read_batch_utf8(10, paths.clone()); + futures::pin_mut!(stream); + + let item1 = stream.next().await.unwrap(); + assert_eq!(item1.0, paths[0]); + + let item2 = stream.next().await.unwrap(); + assert_eq!(item2.0, paths[1]); + + // No more items + assert!(stream.next().await.is_none()); + } +} diff --git a/crates/forge_infra/src/fs_read_dir.rs b/crates/forge_infra/src/fs_read_dir.rs new file mode 100644 index 0000000000000000000000000000000000000000..761103fe71e831f4140dd0bc991ffa5c4080570f --- /dev/null +++ b/crates/forge_infra/src/fs_read_dir.rs @@ -0,0 +1,276 @@ +use std::path::{Path, PathBuf}; + +use anyhow::Result; +use forge_app::DirectoryReaderInfra; +use forge_fs::ForgeFS; +use forge_walker::Walker; +use futures::StreamExt; +use glob::Pattern; + +/// Service for reading multiple files from a directory asynchronously +pub struct ForgeDirectoryReaderService { + parallel_file_reads: usize, +} + +impl ForgeDirectoryReaderService { + /// Creates a new service with the given concurrency cap for parallel reads. + /// + /// # Arguments + /// + /// * `parallel_file_reads` - Maximum number of files to read concurrently + pub fn new(parallel_file_reads: usize) -> Self { + Self { parallel_file_reads } + } + + /// Lists all entries in a directory without reading file contents + /// Returns a vector of tuples containing (entry_path, is_directory) + /// Much more efficient than read_directory_files for directory listings + /// Respects .gitignore, .ignore, and other standard ignore files + async fn list_directory_entries(&self, directory: &Path) -> Result> { + // Check if directory exists + if !ForgeFS::exists(directory) || ForgeFS::is_file(directory) { + return Ok(vec![]); + } + + // Use Walker to get entries with gitignore filtering + let files = Walker::max_all() + .cwd(directory.to_path_buf()) + .max_depth(1) + .skip_binary(true) + .hidden(true) + .get() + .await?; + + let mut entries = Vec::new(); + + for file in files { + // Skip root directory entry + if file.path == "/" { + continue; + } + + let file_path = PathBuf::from(file.path); + let absolute_path = if file_path.is_relative() { + directory.join(file_path) + } else { + file_path + }; + let is_dir = absolute_path.is_dir(); + entries.push((absolute_path, is_dir)); + } + + Ok(entries) + } + + /// Reads all files in a directory that match the given filter pattern + /// Returns a vector of tuples containing (file_path, file_content) + /// Files are read asynchronously/in parallel for better performance + async fn read_directory_files( + &self, + directory: &Path, + pattern: Option<&str>, + ) -> Result> { + // Check if directory exists + if !ForgeFS::exists(directory) || ForgeFS::is_file(directory) { + return Ok(vec![]); + } + + // Build glob pattern if filter is provided + let glob_pattern = if let Some(pattern) = pattern { + Some(Pattern::new(pattern)?) + } else { + None + }; + + // Read directory entries + let mut dir = ForgeFS::read_dir(directory).await?; + let mut file_paths = Vec::new(); + + while let Some(entry) = dir.next_entry().await? { + let path = entry.path(); + + // Only process files (not directories) + if ForgeFS::is_file(&path) { + // Apply filter if provided + if let Some(ref pattern) = glob_pattern { + if let Some(file_name) = path.file_name().and_then(|n| n.to_str()) + && pattern.matches(file_name) + { + file_paths.push(path); + } + } else { + file_paths.push(path); + } + } + } + + // Read files in parallel with a concurrency cap to avoid EMFILE errors + let files = futures::stream::iter(file_paths) + .map(|path| async move { + let path_clone = path.clone(); + match ForgeFS::read_to_string(&path).await { + Ok(content) => Some((path_clone, content)), + Err(_) => None, // Skip files that can't be read + } + }) + .buffer_unordered(self.parallel_file_reads) + .filter_map(std::future::ready) + .collect::>() + .await; + + Ok(files) + } +} + +#[async_trait::async_trait] +impl DirectoryReaderInfra for ForgeDirectoryReaderService { + async fn list_directory_entries(&self, directory: &Path) -> Result> { + self.list_directory_entries(directory).await + } + + async fn read_directory_files( + &self, + directory: &Path, + pattern: Option<&str>, + ) -> Result> { + self.read_directory_files(directory, pattern).await + } +} + +#[cfg(test)] +mod tests { + use std::fs; + + use pretty_assertions::assert_eq; + use tempfile::tempdir; + + use super::*; + + fn write_file(path: &Path, content: &str) { + fs::write(path, content).unwrap(); + } + + #[tokio::test] + async fn test_read_directory_files_with_filter() { + let fixture = tempdir().unwrap(); + write_file(&fixture.path().join("test.md"), "# Markdown content"); + write_file(&fixture.path().join("test.txt"), "Text content"); + write_file(&fixture.path().join("test.rs"), "fn main() {}"); + + let actual = ForgeDirectoryReaderService::new(64) + .read_directory_files(fixture.path(), Some("*.md")) + .await + .unwrap(); + + let expected = vec![( + fixture.path().join("test.md"), + "# Markdown content".to_string(), + )]; + assert_eq!(actual, expected); + } + + #[tokio::test] + async fn test_read_directory_files_without_filter() { + let fixture = tempdir().unwrap(); + write_file(&fixture.path().join("file1.txt"), "Content 1"); + write_file(&fixture.path().join("file2.md"), "Content 2"); + + let mut actual = ForgeDirectoryReaderService::new(64) + .read_directory_files(fixture.path(), None) + .await + .unwrap(); + actual.sort_by(|(a, _), (b, _)| a.file_name().cmp(&b.file_name())); + + let expected = vec![ + (fixture.path().join("file1.txt"), "Content 1".to_string()), + (fixture.path().join("file2.md"), "Content 2".to_string()), + ]; + assert_eq!(actual, expected); + } + + #[tokio::test] + async fn test_read_directory_files_nonexistent_directory() { + let actual = ForgeDirectoryReaderService::new(64) + .read_directory_files(Path::new("/nonexistent"), None) + .await + .unwrap(); + + let expected: Vec<(PathBuf, String)> = vec![]; + assert_eq!(actual, expected); + } + + #[tokio::test] + async fn test_read_directory_files_ignores_subdirectories() { + let fixture = tempdir().unwrap(); + write_file(&fixture.path().join("test.txt"), "File content"); + + let subdir = fixture.path().join("subdir"); + fs::create_dir(&subdir).unwrap(); + write_file(&subdir.join("subfile.txt"), "Sub content"); + + let actual = ForgeDirectoryReaderService::new(64) + .read_directory_files(fixture.path(), None) + .await + .unwrap(); + + let expected = vec![(fixture.path().join("test.txt"), "File content".to_string())]; + assert_eq!(actual, expected); + } + + #[tokio::test] + async fn test_list_directory_entries() { + let fixture = tempdir().unwrap(); + write_file(&fixture.path().join("file1.txt"), "Content 1"); + write_file(&fixture.path().join("file2.md"), "Content 2"); + + let subdir = fixture.path().join("subdir"); + fs::create_dir(&subdir).unwrap(); + + let mut actual = ForgeDirectoryReaderService::new(64) + .list_directory_entries(fixture.path()) + .await + .unwrap(); + actual.sort_by(|(a, _), (b, _)| a.file_name().cmp(&b.file_name())); + + let expected = vec![ + (fixture.path().join("file1.txt"), false), + (fixture.path().join("file2.md"), false), + (fixture.path().join("subdir"), true), + ]; + assert_eq!(actual, expected); + } + + #[tokio::test] + async fn test_list_directory_entries_nonexistent() { + let actual = ForgeDirectoryReaderService::new(64) + .list_directory_entries(Path::new("/nonexistent")) + .await + .unwrap(); + + let expected: Vec<(PathBuf, bool)> = vec![]; + assert_eq!(actual, expected); + } + + #[tokio::test] + async fn test_list_directory_entries_respects_gitignore() { + let fixture = tempdir().unwrap(); + write_file(&fixture.path().join("file.txt"), "Content"); + write_file(&fixture.path().join("ignored.log"), "Log"); + + let subdir = fixture.path().join("node_modules"); + fs::create_dir(&subdir).unwrap(); + + let git_dir = fixture.path().join(".git"); + fs::create_dir(&git_dir).unwrap(); + + write_file(&fixture.path().join(".gitignore"), "*.log\nnode_modules/\n"); + + let actual = ForgeDirectoryReaderService::new(64) + .list_directory_entries(fixture.path()) + .await + .unwrap(); + + let expected = vec![(fixture.path().join("file.txt"), false)]; + assert_eq!(actual, expected); + } +} diff --git a/crates/forge_infra/src/mcp_client.rs b/crates/forge_infra/src/mcp_client.rs new file mode 100644 index 0000000000000000000000000000000000000000..1c0b5db47ac2af12ce47cd858a7efe2b99e0113f --- /dev/null +++ b/crates/forge_infra/src/mcp_client.rs @@ -0,0 +1,832 @@ +use std::borrow::Cow; +use std::collections::BTreeMap; +use std::future::Future; +use std::sync::{Arc, OnceLock, RwLock}; + +use backon::{ExponentialBuilder, Retryable}; +use bstr::ByteSlice; +use forge_app::McpClientInfra; +use forge_domain::{ + Environment, Image, McpHttpServer, McpServerConfig, ToolDefinition, ToolName, ToolOutput, +}; +use http::{HeaderName, HeaderValue}; +use rmcp::model::{CallToolRequestParams, ClientInfo, Implementation, InitializeRequestParams}; +use rmcp::service::RunningService; +use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig; +use rmcp::transport::{StreamableHttpClientTransport, TokioChildProcess}; +use rmcp::{RoleClient, ServiceExt}; +use schemars::Schema; +use serde_json::Value; +use tokio::io::{AsyncBufReadExt, BufReader}; +use tokio::process::Command; + +use crate::error::Error; + +const VERSION: &str = match option_env!("APP_VERSION") { + Some(val) => val, + None => env!("CARGO_PKG_VERSION"), +}; + +type RmcpClient = RunningService; + +#[derive(Clone)] +pub struct ForgeMcpClient { + client: Arc>>>, + config: McpServerConfig, + env_vars: BTreeMap, + environment: Environment, + resolved_config: Arc>>, +} + +impl ForgeMcpClient { + pub fn new( + config: McpServerConfig, + env_vars: &BTreeMap, + environment: Environment, + ) -> Self { + Self { + client: Default::default(), + config, + env_vars: env_vars.clone(), + environment, + resolved_config: Arc::new(OnceLock::new()), + } + } + + /// Gets the resolved configuration, lazily initializing templates if needed + fn get_resolved_config(&self) -> anyhow::Result<&McpServerConfig> { + self.resolved_config + .get_or_init(|| match &self.config { + McpServerConfig::Http(http) => { + resolve_http_templates(http.clone(), &self.env_vars).map(McpServerConfig::Http) + } + x => Ok(x.clone()), + }) + .as_ref() + .map_err(|e| anyhow::anyhow!("{e}")) + } + + fn client_info(&self) -> ClientInfo { + ClientInfo::new(Default::default(), Implementation::new("Forge", VERSION)) + } + + /// Connects to the MCP server. If `force` is true, it will reconnect even + /// if already connected. + async fn connect(&self) -> anyhow::Result> { + if let Some(client) = self.get_client() { + Ok(client.clone()) + } else { + let client = self.create_connection().await?; + self.set_client(client.clone()); + Ok(client.clone()) + } + } + + fn get_client(&self) -> Option> { + self.client.read().ok().and_then(|guard| guard.clone()) + } + + fn set_client(&self, client: Arc) { + if let Ok(mut guard) = self.client.write() { + *guard = Some(client); + } + } + + async fn create_connection(&self) -> anyhow::Result> { + let config = self.get_resolved_config()?; + let client = match config { + McpServerConfig::Stdio(stdio) => { + let mut cmd = Command::new(stdio.command.clone()); + + for (key, value) in &stdio.env { + cmd.env(key, value); + } + + cmd.args(&stdio.args).kill_on_drop(true); + + // Use builder pattern to capture stderr + let (transport, stderr) = TokioChildProcess::builder(cmd) + .stderr(std::process::Stdio::piped()) + .spawn()?; + + // Spawn a task to drain stderr to prevent buffer overflow + // If stderr fills up, the child process will block + if let Some(stderr) = stderr { + tokio::spawn(async move { + let mut reader = BufReader::new(stderr).lines(); + while let Ok(Some(line)) = reader.next_line().await { + tracing::warn!("MCP server stderr: {}", line); + } + }); + } + Arc::new(self.client_info().serve(transport).await?) + } + McpServerConfig::Http(http) => { + // Check if OAuth is explicitly disabled + if http.is_oauth_disabled() { + // OAuth explicitly disabled - only try standard connection + Arc::new(self.create_standard_http_connection(http).await?) + } else if let Some(oauth_config) = http.oauth_config() { + // OAuth explicitly configured - use it directly + // Do NOT allow interactive auth during normal connection + self.create_oauth_connection(http, oauth_config, false) + .await? + } else { + // Auto-detect: try standard first, fall back to OAuth on auth errors + match self.create_standard_http_connection(http).await { + Ok(client) => Arc::new(client), + Err(e) => { + let error_str = e.to_string().to_lowercase(); + if error_str.contains("401") + || error_str.contains("unauthorized") + || error_str.contains("authentication required") + || error_str.contains("auth required") + || error_str.contains("oauth") + { + tracing::info!( + "Standard connection failed with auth error for: {}, trying stored credentials", + http.url + ); + // Try OAuth with stored credentials (non-interactive) + // If stored credentials exist, use them; otherwise error + let default_config = forge_domain::McpOAuthConfig::default(); + self.create_oauth_connection(http, &default_config, false) + .await? + } else { + return Err(e); + } + } + } + } + } + }; + + Ok(client) + } + + /// Create a standard HTTP connection without OAuth + async fn create_standard_http_connection( + &self, + http: &McpHttpServer, + ) -> anyhow::Result { + let config = StreamableHttpClientTransportConfig::with_uri(http.url.clone()) + .custom_headers(build_header_map(&http.headers)); + let transport = StreamableHttpClientTransport::from_config(config); + Ok(self.client_info().serve(transport).await?) + } + + /// Create an OAuth-enabled connection using rmcp's OAuth support. + /// + /// Uses rmcp's `AuthorizationManager` and `OAuthState` state machine which + /// properly handle: + /// 1. OAuth metadata discovery via RFC 8414 + /// 2. Dynamic client registration via RFC 7591 + /// 3. PKCE challenge/verifier generation and validation + /// 4. CSRF state parameter generation and validation + /// 5. Authorization code exchange for tokens + /// 6. Token refresh via refresh_token grant + /// 7. Token persistence via `CredentialStore` trait + /// + /// # Arguments + /// * `allow_interactive` - If true, will open browser for user + /// authentication if no stored credentials exist. If false, returns an + /// error instead. + async fn create_oauth_connection( + &self, + http: &McpHttpServer, + oauth_config: &forge_domain::McpOAuthConfig, + allow_interactive: bool, + ) -> anyhow::Result> { + use rmcp::transport::auth::{AuthorizationManager, OAuthState}; + + use crate::auth::McpTokenStorage; + + let credential_store = McpTokenStorage::new(http.url.clone(), self.environment.clone()); + + // First, try to use cached credentials with auto-refresh + let mut auth_manager = AuthorizationManager::new(&http.url) + .await + .map_err(|e| anyhow::anyhow!("Failed to create OAuth manager: {}", e))?; + + auth_manager.set_credential_store(credential_store); + + // Try to load and use stored credentials (with automatic token refresh) + match auth_manager.initialize_from_store().await { + Ok(true) => { + // Stored credentials loaded. Try to get a valid access token + // (this auto-refreshes if expired and refresh_token is available) + match auth_manager.get_access_token().await { + Ok(token) => { + tracing::debug!("Using stored/refreshed OAuth token for: {}", http.url); + return self.connect_with_token(http, &token).await; + } + Err(e) => { + tracing::warn!( + "Stored token invalid for {}: {}, re-authenticating", + http.url, + e + ); + } + } + } + Ok(false) => { + tracing::info!("No stored credentials for: {}", http.url); + } + Err(e) => { + tracing::warn!("Failed to load stored credentials for {}: {}", http.url, e); + } + } + + // No valid cached credentials + if !allow_interactive { + // Interactive auth not allowed - return error with instructions + return Err(anyhow::anyhow!( + "MCP server '{}' requires authentication. Run 'mcp login ' to authenticate.", + http.url + )); + } + + // Interactive auth allowed - start full OAuth authorization flow + // Create a fresh OAuthState to run the browser-based flow + let mut oauth_state = OAuthState::new(&http.url, None) + .await + .map_err(|e| anyhow::anyhow!("Failed to initialize OAuth state: {}", e))?; + + let redirect_uri = oauth_config + .redirect_uri + .clone() + .unwrap_or_else(|| "http://127.0.0.1:8765/callback".to_string()); + + let scopes: Vec<&str> = oauth_config.scopes.iter().map(|s| s.as_str()).collect(); + + // start_authorization discovers metadata, registers client, generates PKCE + + // CSRF state + oauth_state + .start_authorization(&scopes, &redirect_uri, Some("Forge")) + .await + .map_err(|e| anyhow::anyhow!("OAuth authorization flow failed: {}", e))?; + + // Get the authorization URL (includes PKCE challenge and CSRF state) + let auth_url = oauth_state + .get_authorization_url() + .await + .map_err(|e| anyhow::anyhow!("Failed to get authorization URL: {}", e))?; + + tracing::info!("Starting OAuth authentication for MCP server: {}", http.url); + + // Parse redirect URI to get port for callback server + let redirect_url: url::Url = redirect_uri + .parse() + .map_err(|e| anyhow::anyhow!("Invalid redirect URI: {}", e))?; + let port = redirect_url.port().unwrap_or(8765); + + // Start local callback server, open browser, wait for redirect + let (code, state) = self.run_oauth_callback_server(port, &auth_url).await?; + + // Exchange authorization code for tokens (validates CSRF state internally) + // rmcp's OAuthState handles PKCE verifier inclusion in the token request + oauth_state + .handle_callback(&code, &state) + .await + .map_err(|e| anyhow::anyhow!("Failed to exchange authorization code: {}", e))?; + + // Get the access token from the completed OAuth flow + let access_token = oauth_state + .get_access_token() + .await + .map_err(|e| anyhow::anyhow!("Failed to get access token after OAuth: {}", e))?; + + // Save credentials for future use via our persistent store + let credentials = oauth_state + .get_credentials() + .await + .map_err(|e| anyhow::anyhow!("Failed to get credentials: {}", e))?; + + { + use rmcp::transport::auth::CredentialStore; + let save_store = McpTokenStorage::new(http.url.clone(), self.environment.clone()); + let stored = rmcp::transport::auth::StoredCredentials::new( + credentials.0, + credentials.1, + vec![], + None, + ); + save_store + .save(stored) + .await + .map_err(|e| anyhow::anyhow!("Failed to save credentials: {}", e))?; + } + + tracing::info!( + "OAuth authentication successful for MCP server: {}", + http.url + ); + + self.connect_with_token(http, &access_token).await + } + + /// Connect to an MCP server using a bearer token. + /// + /// Uses StreamableHTTP transport only - does NOT fall back to SSE + /// since SSE transport doesn't support auth headers in the same way. + /// Auth errors are transport-independent so falling back to SSE + /// with the same auth issue would be pointless. + async fn connect_with_token( + &self, + http: &McpHttpServer, + token: &str, + ) -> anyhow::Result> { + let config = + StreamableHttpClientTransportConfig::with_uri(http.url.clone()).auth_header(token); + let transport = StreamableHttpClientTransport::from_config(config); + Ok(Arc::new(self.client_info().serve(transport).await?)) + } + + /// Runs a local HTTP server to receive the OAuth callback, opens the + /// browser, and returns the authorization code and state. + async fn run_oauth_callback_server( + &self, + port: u16, + auth_url: &str, + ) -> anyhow::Result<(String, String)> { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::TcpListener; + + let listener = TcpListener::bind(format!("127.0.0.1:{}", port)) + .await + .map_err(|e| { + anyhow::anyhow!( + "Failed to start OAuth callback server on port {}: {}. \ + Is another process using this port?", + port, + e + ) + })?; + + tracing::info!("OAuth callback server listening on port {}", port); + + // Open browser + if let Err(e) = open::that(auth_url) { + tracing::warn!( + "Failed to open browser: {}. Please open this URL manually:\n{}", + e, + auth_url + ); + eprintln!( + "\nPlease open this URL in your browser to authenticate:\n{}\n", + auth_url + ); + } else { + eprintln!("\nOpening browser for OAuth authentication...\n"); + } + + // Wait for callback with timeout + let timeout = tokio::time::Duration::from_secs(300); // 5 minutes + let (mut stream, _addr) = tokio::time::timeout(timeout, listener.accept()) + .await + .map_err(|_| anyhow::anyhow!("OAuth callback timed out after 5 minutes"))? + .map_err(|e| anyhow::anyhow!("Failed to accept OAuth callback: {}", e))?; + + // Read the HTTP request + let mut buf = vec![0u8; 4096]; + let n = stream.read(&mut buf).await?; + let request = buf.get(..n).unwrap_or(&[]).to_str_lossy(); + let first_line = request.lines().next().unwrap_or(""); + let path = first_line.split_whitespace().nth(1).unwrap_or("/"); + + // Parse query parameters + let query_start = path.find('?').unwrap_or(path.len()); + let query_string = path.get(query_start..).unwrap_or(""); + let params: std::collections::HashMap = + url::form_urlencoded::parse(query_string.trim_start_matches('?').as_bytes()) + .into_owned() + .collect(); + + let code = params + .get("code") + .ok_or_else(|| { + let error = params.get("error").map(|e| e.as_str()).unwrap_or("unknown"); + let desc = params + .get("error_description") + .map(|d| d.as_str()) + .unwrap_or("No description"); + anyhow::anyhow!("OAuth error: {} - {}", error, desc) + })? + .clone(); + + let state = params + .get("state") + .ok_or_else(|| anyhow::anyhow!("Missing state parameter in OAuth callback"))? + .clone(); + + // Send styled success response with auto-close + let response_body = r#"Forge - Authorization Successful

Authorization Successful

You can close this window and return to Forge.

"#; + let response = format!( + "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", + response_body.len(), + response_body + ); + let _ = stream.write_all(response.as_bytes()).await; + + Ok((code, state)) + } + + async fn list(&self) -> anyhow::Result> { + let client = self.connect().await?; + let tools = client.list_tools(None).await?; + Ok(tools + .tools + .into_iter() + .filter_map(|tool| { + Some( + ToolDefinition::new(tool.name) + .description(tool.description.unwrap_or_default()) + .input_schema( + serde_json::from_value::(Value::Object( + tool.input_schema.as_ref().clone(), + )) + .ok()?, + ), + ) + }) + .collect()) + } + + async fn call(&self, tool_name: &ToolName, input: &Value) -> anyhow::Result { + let client = self.connect().await?; + let result = client + .call_tool({ + let mut params = CallToolRequestParams::new(Cow::Owned(tool_name.to_string())); + if let Value::Object(args) = input { + params = params.with_arguments(args.clone()); + } + params + }) + .await?; + + let tool_contents: Vec = result + .content + .into_iter() + .map(|content| match content.raw { + rmcp::model::RawContent::Text(raw_text_content) => { + Ok(ToolOutput::text(raw_text_content.text)) + } + rmcp::model::RawContent::Image(raw_image_content) => Ok(ToolOutput::image( + Image::new_base64(raw_image_content.data, raw_image_content.mime_type.as_str()), + )), + rmcp::model::RawContent::Resource(_) => { + Err(Error::UnsupportedMcpResponse("Resource").into()) + } + rmcp::model::RawContent::ResourceLink(_) => { + Err(Error::UnsupportedMcpResponse("ResourceLink").into()) + } + rmcp::model::RawContent::Audio(_) => { + Err(Error::UnsupportedMcpResponse("Audio").into()) + } + }) + .collect::>>()?; + + Ok(ToolOutput::from(tool_contents.into_iter()) + .is_error(result.is_error.unwrap_or_default())) + } + + async fn attempt_with_retry(&self, call: impl Fn() -> F) -> anyhow::Result + where + F: Future>, + { + call.retry( + ExponentialBuilder::default() + .with_max_times(5) + .with_jitter(), + ) + .when(|err| { + let is_transport = err + .downcast_ref::() + .map(|e| { + matches!( + e, + rmcp::ServiceError::TransportSend(_) | rmcp::ServiceError::TransportClosed + ) + }) + .unwrap_or(false); + + if is_transport && let Ok(mut guard) = self.client.write() { + guard.take(); + } + + is_transport + }) + .await + } +} + +#[async_trait::async_trait] +impl McpClientInfra for ForgeMcpClient { + async fn list(&self) -> anyhow::Result> { + self.attempt_with_retry(|| self.list()).await + } + + async fn call(&self, tool_name: &ToolName, input: Value) -> anyhow::Result { + self.attempt_with_retry(|| self.call(tool_name, &input)) + .await + } +} + +/// Resolves mustache templates in McpHttpServer headers using Handlebars +/// and provided environment variables +fn resolve_http_templates( + mut http: McpHttpServer, + env_vars: &BTreeMap, +) -> anyhow::Result { + let handlebars = forge_app::TemplateEngine::handlebar_instance(); + + // Create template data with env variables nested under "env" + let template_data = serde_json::json!({"env": env_vars}); + + // Resolve templates in headers + for value in http.headers.values_mut() { + // Try to render the template, but keep original value if it fails + if let Ok(resolved) = handlebars.render_template(value, &template_data) { + *value = resolved; + } + } + + Ok(http) +} + +fn build_header_map( + headers: &BTreeMap, +) -> std::collections::HashMap { + headers + .iter() + .filter_map(|(k, v)| { + let name = k.parse::().ok()?; + let val = v.parse::().ok()?; + Some((name, val)) + }) + .collect() +} + +/// Trigger OAuth authentication for a specific MCP server URL. +/// +/// Runs the full OAuth flow: metadata discovery, dynamic registration, +/// browser-based authorization, and token persistence. +/// +/// # Arguments +/// * `server_url` - The URL of the MCP server to authenticate with +/// * `env` - The environment for file system paths +pub async fn mcp_auth(server_url: &str, env: &Environment) -> anyhow::Result<()> { + use rmcp::transport::auth::{CredentialStore, OAuthState}; + + use crate::auth::McpTokenStorage; + + // Start fresh OAuth flow via OAuthState + let mut oauth_state = OAuthState::new(server_url, None) + .await + .map_err(|e| anyhow::anyhow!("Failed to initialize OAuth state: {}", e))?; + + let redirect_uri = "http://127.0.0.1:8765/callback"; + + oauth_state + .start_authorization(&[], redirect_uri, Some("Forge")) + .await + .map_err(|e| anyhow::anyhow!("OAuth authorization flow failed: {}", e))?; + + let auth_url = oauth_state + .get_authorization_url() + .await + .map_err(|e| anyhow::anyhow!("Failed to get authorization URL: {}", e))?; + + // Start callback server and open browser + let listener = tokio::net::TcpListener::bind("127.0.0.1:8765") + .await + .map_err(|e| anyhow::anyhow!("Failed to start OAuth callback server: {}", e))?; + + if let Err(e) = open::that(&auth_url) { + tracing::warn!("Failed to open browser: {}", e); + eprintln!( + "\nPlease open this URL in your browser to authenticate:\n{}\n", + auth_url + ); + } else { + eprintln!("\nOpening browser for OAuth authentication...\n"); + } + + let timeout = tokio::time::Duration::from_secs(300); + let (mut stream, _) = tokio::time::timeout(timeout, listener.accept()) + .await + .map_err(|_| anyhow::anyhow!("OAuth callback timed out after 5 minutes"))? + .map_err(|e| anyhow::anyhow!("Failed to accept OAuth callback: {}", e))?; + + // Read HTTP request and parse callback params + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let mut buf = vec![0u8; 4096]; + let n = stream.read(&mut buf).await?; + let request = buf.get(..n).unwrap_or(&[]).to_str_lossy(); + let first_line = request.lines().next().unwrap_or(""); + let path = first_line.split_whitespace().nth(1).unwrap_or("/"); + let query_start = path.find('?').unwrap_or(path.len()); + let params: std::collections::HashMap = url::form_urlencoded::parse( + path.get(query_start..) + .unwrap_or("") + .trim_start_matches('?') + .as_bytes(), + ) + .into_owned() + .collect(); + + let code = params + .get("code") + .ok_or_else(|| { + let error = params.get("error").map(|e| e.as_str()).unwrap_or("unknown"); + let desc = params + .get("error_description") + .map(|d| d.as_str()) + .unwrap_or("No description"); + anyhow::anyhow!("OAuth error: {} - {}", error, desc) + })? + .clone(); + + let state = params + .get("state") + .ok_or_else(|| anyhow::anyhow!("Missing state parameter in OAuth callback"))? + .clone(); + + // Send styled response + let body = r#"Forge - Authorization Successful

Authorization Successful

You can close this window and return to Forge.

"#; + let resp = format!( + "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", + body.len(), + body + ); + let _ = stream.write_all(resp.as_bytes()).await; + + // Exchange code for tokens + oauth_state + .handle_callback(&code, &state) + .await + .map_err(|e| anyhow::anyhow!("Failed to exchange authorization code: {}", e))?; + + // Save credentials + let credentials = oauth_state + .get_credentials() + .await + .map_err(|e| anyhow::anyhow!("Failed to get credentials: {}", e))?; + + let save_store = McpTokenStorage::new(server_url.to_string(), env.clone()); + let stored = + rmcp::transport::auth::StoredCredentials::new(credentials.0, credentials.1, vec![], None); + save_store + .save(stored) + .await + .map_err(|e| anyhow::anyhow!("Failed to save credentials: {}", e))?; + + Ok(()) +} + +/// Remove stored OAuth credentials for a specific MCP server. +/// +/// # Arguments +/// * `server_url` - The URL of the MCP server to remove credentials for +/// * `env` - The environment for file system paths +pub async fn mcp_logout(server_url: &str, env: &Environment) -> anyhow::Result<()> { + use crate::auth::McpTokenStorage; + let storage = McpTokenStorage::new(server_url.to_string(), env.clone()); + storage.remove_credentials().await +} + +/// Remove all stored MCP OAuth credentials. +/// +/// # Arguments +/// * `env` - The environment for file system paths +pub async fn mcp_logout_all(env: &Environment) -> anyhow::Result<()> { + use crate::auth::McpCredentialStore; + let path = McpCredentialStore::credential_path(env); + if path.exists() { + tokio::fs::remove_file(&path).await?; + } + Ok(()) +} + +/// Get the auth status for a specific MCP server. +/// +/// Returns one of: "authenticated", "expired", "not_authenticated" +/// +/// # Arguments +/// * `server_url` - The URL of the MCP server +/// * `env` - The environment for file system paths +pub async fn mcp_auth_status(server_url: &str, env: &Environment) -> String { + use crate::auth::McpTokenStorage; + let storage = McpTokenStorage::new(server_url.to_string(), env.clone()); + match storage.load_credentials().await { + Ok(Some(entry)) => { + if let Some(expires_at) = entry.tokens.expires_at { + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); + if expires_at <= now { + if entry.tokens.refresh_token.is_some() { + "expired (has refresh token)".to_string() + } else { + "expired".to_string() + } + } else { + "authenticated".to_string() + } + } else { + "authenticated".to_string() + } + } + Ok(None) => "not authenticated".to_string(), + Err(_) => "unknown (error reading credentials)".to_string(), + } +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::*; + + #[test] + fn test_resolve_http_templates_with_env() { + let env_vars = BTreeMap::from([ + ("GH_TOKEN".to_string(), "secret_token_123".to_string()), + ("API_KEY".to_string(), "api_key_456".to_string()), + ]); + + let http = McpHttpServer { + url: "https://api.example.com".to_string(), + headers: BTreeMap::from([ + ( + "Authorization".to_string(), + "Bearer {{env.GH_TOKEN}}".to_string(), + ), + ("X-API-Key".to_string(), "{{env.API_KEY}}".to_string()), + ("Content-Type".to_string(), "application/json".to_string()), + ]), + timeout: None, + disable: false, + oauth: Default::default(), + }; + + let resolved = resolve_http_templates(http, &env_vars).unwrap(); + + assert_eq!( + resolved.headers.get("Authorization"), + Some(&"Bearer secret_token_123".to_string()) + ); + assert_eq!( + resolved.headers.get("X-API-Key"), + Some(&"api_key_456".to_string()) + ); + assert_eq!( + resolved.headers.get("Content-Type"), + Some(&"application/json".to_string()) + ); + } + + #[test] + fn test_resolve_http_templates_missing_env_var() { + let env_vars = BTreeMap::new(); // Empty env vars + + let http = McpHttpServer { + url: "https://api.example.com".to_string(), + headers: BTreeMap::from([( + "Authorization".to_string(), + "Bearer {{env.MISSING_VAR}}".to_string(), + )]), + timeout: None, + disable: false, + oauth: Default::default(), + }; + + let resolved = resolve_http_templates(http, &env_vars).unwrap(); + + // Should keep original value if template rendering fails + assert_eq!( + resolved.headers.get("Authorization"), + Some(&"Bearer {{env.MISSING_VAR}}".to_string()) + ); + } + + #[test] + fn test_resolve_http_templates_preserves_url_and_disable() { + let env_vars = BTreeMap::from([("TOKEN".to_string(), "test".to_string())]); + + let http = McpHttpServer { + url: "https://test.example.com".to_string(), + headers: BTreeMap::from([("Auth".to_string(), "{{env.TOKEN}}".to_string())]), + timeout: None, + disable: true, + oauth: Default::default(), + }; + + let resolved = resolve_http_templates(http, &env_vars).unwrap(); + + assert_eq!(resolved.url, "https://test.example.com"); + assert_eq!(resolved.disable, true); + assert_eq!(resolved.headers.get("Auth"), Some(&"test".to_string())); + } +} diff --git a/crates/forge_json_repair/src/error.rs b/crates/forge_json_repair/src/error.rs new file mode 100644 index 0000000000000000000000000000000000000000..777be67726a05c85bb889487646579a6ecc4ef39 --- /dev/null +++ b/crates/forge_json_repair/src/error.rs @@ -0,0 +1,27 @@ +use thiserror::Error; + +#[derive(Error, Debug)] +pub enum JsonRepairError { + #[error("Invalid character {character:?} at position {position}")] + InvalidCharacter { character: char, position: usize }, + + #[error("Unexpected character {character:?} at position {position}")] + UnexpectedCharacter { character: char, position: usize }, + + #[error("Unexpected end of JSON string at position {position}")] + UnexpectedEnd { position: usize }, + + #[error("Object key expected at position {position}")] + ObjectKeyExpected { position: usize }, + + #[error("Colon expected at position {position}")] + ColonExpected { position: usize }, + + #[error("Invalid unicode character {chars:?} at position {position}")] + InvalidUnicodeCharacter { chars: String, position: usize }, + + #[error("JSON parsing error: {0}")] + JsonError(#[from] serde_json::Error), +} + +pub type Result = std::result::Result; diff --git a/crates/forge_json_repair/src/lib.rs b/crates/forge_json_repair/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..03af4fd678267ceeb7b836cee88d2abd0c57691f --- /dev/null +++ b/crates/forge_json_repair/src/lib.rs @@ -0,0 +1,7 @@ +mod error; +mod parser; +mod schema_coercion; + +pub use error::{JsonRepairError, Result}; +pub use parser::json_repair; +pub use schema_coercion::coerce_to_schema; diff --git a/crates/forge_json_repair/src/parser.rs b/crates/forge_json_repair/src/parser.rs new file mode 100644 index 0000000000000000000000000000000000000000..35cc652b03df56cc4b600a3c059af919c03bff1e --- /dev/null +++ b/crates/forge_json_repair/src/parser.rs @@ -0,0 +1,1073 @@ +use serde::Deserialize; + +use crate::error::{JsonRepairError, Result}; + +pub struct JsonRepairParser { + chars: Vec, + i: usize, + output: String, +} + +impl JsonRepairParser { + pub fn new(text: String) -> Self { + Self { chars: text.chars().collect(), i: 0, output: String::new() } + } + + pub fn parse Deserialize<'de>>(mut self) -> Result { + self.parse_markdown_code_block(&["```", "[```", "{```"]); + + let processed = self.parse_value()?; + if !processed { + return Err(JsonRepairError::UnexpectedEnd { position: self.chars.len() }); + } + + self.parse_markdown_code_block(&["```", "```]", "```}"]); + + let processed_comma = self.parse_character(','); + if processed_comma { + self.parse_whitespace_and_skip_comments(true); + } + + if self.is_start_of_value(self.current_char()) && self.ends_with_comma_or_newline() { + if !processed_comma { + self.output = self.insert_before_last_whitespace(","); + } + self.parse_newline_delimited_json()?; + } else if processed_comma { + self.output = self.strip_last_occurrence(","); + } + + // repair redundant end quotes + while matches!(self.current_char(), Some('}' | ']')) { + self.i += 1; + self.parse_whitespace_and_skip_comments(true); + } + + if self.i >= self.chars.len() { + return Ok(serde_json::from_str::(&self.output)?); + } + + Err(JsonRepairError::UnexpectedCharacter { + character: self.current_char().unwrap_or('\0'), + position: self.i, + }) + } + + fn current_char(&self) -> Option { + self.chars.get(self.i).copied() + } + + fn parse_value(&mut self) -> Result { + self.parse_whitespace_and_skip_comments(true); + + let processed = self.parse_object()? + || self.parse_array()? + || self.parse_string(false, None)? + || self.parse_number()? + || self.parse_keywords() + || self.parse_unquoted_string(false)? + || self.parse_regex()?; + + self.parse_whitespace_and_skip_comments(true); + Ok(processed) + } + + fn parse_whitespace_and_skip_comments(&mut self, skip_newline: bool) -> bool { + let start = self.i; + + self.parse_whitespace(skip_newline); + loop { + if self.parse_comment() { + self.parse_whitespace(skip_newline); + } else { + break; + } + } + + self.i > start + } + + fn parse_whitespace(&mut self, skip_newline: bool) -> bool { + let mut whitespace = String::new(); + + while let Some(ch) = self.current_char() { + if self.is_whitespace(ch, skip_newline) { + whitespace.push(ch); + self.i += 1; + } else if self.is_special_whitespace(ch) { + whitespace.push(' '); + self.i += 1; + } else { + break; + } + } + + if !whitespace.is_empty() { + self.output.push_str(&whitespace); + true + } else { + false + } + } + + fn parse_comment(&mut self) -> bool { + // block comment /* ... */ + if self.current_char() == Some('/') && self.chars.get(self.i + 1) == Some(&'*') { + while self.i < self.chars.len() && !self.at_end_of_block_comment() { + self.i += 1; + } + self.i += 2; + return true; + } + + // line comment // ... + if self.current_char() == Some('/') && self.chars.get(self.i + 1) == Some(&'/') { + while self.i < self.chars.len() && self.current_char() != Some('\n') { + self.i += 1; + } + return true; + } + + false + } + + fn parse_markdown_code_block(&mut self, blocks: &[&str]) -> bool { + if self.skip_markdown_code_block(blocks) { + if self.is_function_name_char_start(self.current_char()) { + while self.i < self.chars.len() && self.is_function_name_char(self.current_char()) { + self.i += 1; + } + } + self.parse_whitespace_and_skip_comments(true); + return true; + } + false + } + + fn skip_markdown_code_block(&mut self, blocks: &[&str]) -> bool { + for block in blocks { + let block_chars: Vec = block.chars().collect(); + if self.i + block_chars.len() <= self.chars.len() { + let slice = &self.chars.get(self.i..self.i + block_chars.len()); + if slice == &Some(block_chars.as_slice()) { + self.i += block_chars.len(); + return true; + } + } + } + false + } + + fn parse_character(&mut self, expected: char) -> bool { + if self.current_char() == Some(expected) { + self.output.push(expected); + self.i += 1; + true + } else { + false + } + } + + fn skip_character(&mut self, expected: char) -> bool { + if self.current_char() == Some(expected) { + self.i += 1; + true + } else { + false + } + } + + fn skip_ellipsis(&mut self) -> bool { + self.parse_whitespace_and_skip_comments(true); + + if self.current_char() == Some('.') + && self.chars.get(self.i + 1) == Some(&'.') + && self.chars.get(self.i + 2) == Some(&'.') + { + self.i += 3; + self.parse_whitespace_and_skip_comments(true); + self.skip_character(','); + return true; + } + false + } + + fn parse_object(&mut self) -> Result { + if self.current_char() != Some('{') { + return Ok(false); + } + + self.output.push('{'); + self.i += 1; + self.parse_whitespace_and_skip_comments(true); + + // repair: skip leading comma like in {, message: "hi"} + if self.skip_character(',') { + self.parse_whitespace_and_skip_comments(true); + } + + let mut initial = true; + while self.i < self.chars.len() && self.current_char() != Some('}') { + let _processed_comma = if !initial { + let comma = self.parse_character(','); + if !comma { + // repair missing comma + self.output = self.insert_before_last_whitespace(","); + } + self.parse_whitespace_and_skip_comments(true); + comma + } else { + initial = false; + true + }; + + self.skip_ellipsis(); + + let processed_key = + self.parse_string(false, None)? || self.parse_unquoted_string(true)?; + if !processed_key { + if matches!(self.current_char(), Some('}' | '{' | ']' | '[') | None) { + // repair trailing comma + self.output = self.strip_last_occurrence(","); + } else { + return Err(JsonRepairError::ObjectKeyExpected { position: self.i }); + } + break; + } + + self.parse_whitespace_and_skip_comments(true); + let processed_colon = self.parse_character(':'); + let truncated_text = self.i >= self.chars.len(); + + if !processed_colon { + if self.is_start_of_value(self.current_char()) || truncated_text { + // repair missing colon + self.output = self.insert_before_last_whitespace(":"); + } else { + return Err(JsonRepairError::ColonExpected { position: self.i }); + } + } + + let processed_value = self.parse_value()?; + if !processed_value { + if processed_colon || truncated_text { + // repair missing object value + self.output.push_str("null"); + } else { + return Err(JsonRepairError::ColonExpected { position: self.i }); + } + } + } + + if self.current_char() == Some('}') { + self.output.push('}'); + self.i += 1; + } else { + // repair missing end bracket + self.output = self.insert_before_last_whitespace("}"); + } + + Ok(true) + } + + fn parse_array(&mut self) -> Result { + if self.current_char() != Some('[') { + return Ok(false); + } + + self.output.push('['); + self.i += 1; + self.parse_whitespace_and_skip_comments(true); + + // repair: skip leading comma like in [,1,2,3] + if self.skip_character(',') { + self.parse_whitespace_and_skip_comments(true); + } + + let mut initial = true; + while self.i < self.chars.len() && self.current_char() != Some(']') { + if !initial { + let processed_comma = self.parse_character(','); + if !processed_comma { + // repair missing comma + self.output = self.insert_before_last_whitespace(","); + } + } else { + initial = false; + } + + self.skip_ellipsis(); + + let processed_value = self.parse_value()?; + if !processed_value { + // repair trailing comma + self.output = self.strip_last_occurrence(","); + break; + } + } + + if self.current_char() == Some(']') { + self.output.push(']'); + self.i += 1; + } else { + // repair missing closing array bracket + self.output = self.insert_before_last_whitespace("]"); + } + + Ok(true) + } + + fn parse_newline_delimited_json(&mut self) -> Result<()> { + let mut initial = true; + let mut processed_value = true; + + while processed_value { + if !initial { + let processed_comma = self.parse_character(','); + if !processed_comma { + // repair: add missing comma + self.output = self.insert_before_last_whitespace(","); + } + } else { + initial = false; + } + + processed_value = self.parse_value()?; + } + + if !processed_value { + // repair: remove trailing comma + self.output = self.strip_last_occurrence(","); + } + + // repair: wrap the output inside array brackets + self.output = format!("[\n{}\n]", self.output); + Ok(()) + } + + fn parse_string( + &mut self, + stop_at_delimiter: bool, + stop_at_index: Option, + ) -> Result { + let skip_escape_chars = self.current_char() == Some('\\'); + if skip_escape_chars { + // repair: remove the first escape character + self.i += 1; + } + + if !self.is_quote(self.current_char()) { + return Ok(false); + } + + let quote_char = self.current_char().unwrap(); + let is_end_quote = self.get_end_quote_matcher(quote_char); + + let i_before = self.i; + let o_before = self.output.len(); + + let mut str_content = String::from("\""); + self.i += 1; + + loop { + if self.i >= self.chars.len() { + // end of text, missing end quote + let i_prev = self.prev_non_whitespace_index(self.i - 1); + if !stop_at_delimiter && self.is_delimiter(self.chars.get(i_prev).copied()) { + // retry parsing the string, stopping at the first next delimiter + self.i = i_before; + self.output.truncate(o_before); + return self.parse_string(true, None); + } + + // repair missing quote + str_content = self.insert_before_last_whitespace_str(&str_content, "\""); + self.output.push_str(&str_content); + return Ok(true); + } + + if let Some(stop_idx) = stop_at_index + && self.i == stop_idx + { + str_content = self.insert_before_last_whitespace_str(&str_content, "\""); + self.output.push_str(&str_content); + return Ok(true); + } + + let ch = self.current_char().unwrap(); + + if is_end_quote(ch) { + // end quote - verify if it's legit + let i_quote = self.i; + let o_quote = str_content.len(); + str_content.push('"'); + self.i += 1; + self.output.push_str(&str_content); + + self.parse_whitespace_and_skip_comments(false); + + if stop_at_delimiter + || self.i >= self.chars.len() + || self.is_delimiter(self.current_char()) + || self.is_quote(self.current_char()) + || self.is_digit(self.current_char()) + { + // legitimate end quote + self.parse_concatenated_string()?; + return Ok(true); + } + + let i_prev_char = self.prev_non_whitespace_index(i_quote - 1); + let prev_char = self.chars.get(i_prev_char).copied(); + + if prev_char == Some(',') { + // comma followed by quote - missing end quote before comma + self.i = i_before; + self.output.truncate(o_before); + return self.parse_string(false, Some(i_prev_char)); + } + + if self.is_delimiter(prev_char) { + // not the right end quote - retry with stop at delimiter + self.i = i_before; + self.output.truncate(o_before); + return self.parse_string(true, None); + } + + // revert and continue - unescaped quote + self.output.truncate(o_before); + self.i = i_quote + 1; + let prefix = str_content.get(..o_quote).unwrap_or(""); + str_content = format!("{}\\\"", prefix); + } else if stop_at_delimiter && self.is_unquoted_string_delimiter(ch) { + // stop at delimiter mode + + // test for URL like "https://..." + let url_candidate = if self.chars.get(self.i - 1) == Some(&':') + && i_before + 1 < self.chars.len() + && self.i + 2 <= self.chars.len() + { + self.chars.get(i_before + 1..self.i + 2) + } else { + None + }; + if url_candidate.is_some_and(|s| self.is_url_start(s)) { + while self.i < self.chars.len() + && self.is_url_char(self.current_char().unwrap()) + { + str_content.push(self.current_char().unwrap()); + self.i += 1; + } + } + + // repair missing quote + str_content = self.insert_before_last_whitespace_str(&str_content, "\""); + self.output.push_str(&str_content); + self.parse_concatenated_string()?; + return Ok(true); + } else if ch == '\\' { + // handle escaped content + if let Some(next_ch) = self.chars.get(self.i + 1) { + match next_ch { + '"' | '\\' | '/' | 'b' | 'f' | 'n' | 'r' | 't' => { + let escaped: String = self + .chars + .get(self.i..self.i + 2) + .map(|s| s.iter().collect()) + .unwrap_or_default(); + str_content.push_str(&escaped); + self.i += 2; + } + ',' if skip_escape_chars => { + // Special case: escaped comma in escaped string should be treated as + // delimiter This creates a new array element + str_content = + self.insert_before_last_whitespace_str(&str_content, "\""); + self.output.push_str(&str_content); + self.output.push(','); + self.i += 2; // skip \, + self.parse_whitespace_and_skip_comments(true); + return Ok(true); + } + 'u' => { + let mut j = 2; + while j < 6 && self.i + j < self.chars.len() { + let hex_ch = self.chars.get(self.i + j).copied().unwrap_or('\0'); + if !self.is_hex(hex_ch) { + break; + } + j += 1; + } + + if j == 6 { + let unicode_seq: String = self + .chars + .get(self.i..self.i + 6) + .map(|s| s.iter().collect()) + .unwrap_or_default(); + str_content.push_str(&unicode_seq); + self.i += 6; + } else if self.i + j >= self.chars.len() { + // repair invalid unicode at end + self.i = self.chars.len(); + } else { + // SAFETY: j <= 6 and self.i + j < self.chars.len() (else branch + // above handles >=) + let invalid_chars: String = self + .chars + .get(self.i..self.i + j.min(6)) + .map(|s| s.iter().collect()) + .unwrap_or_default(); + return Err(JsonRepairError::InvalidUnicodeCharacter { + chars: invalid_chars, + position: self.i, + }); + } + } + _ => { + // repair invalid escape character + str_content.push(*next_ch); + self.i += 2; + } + } + } else { + self.i += 1; + } + } else { + // handle regular characters + if ch == '"' && self.chars.get(self.i - 1) != Some(&'\\') { + // repair unescaped double quote + str_content.push_str("\\\""); + self.i += 1; + } else if self.is_control_character(ch) { + // unescaped control character + str_content.push_str(&self.get_control_character_escape(ch)); + self.i += 1; + } else { + if !self.is_valid_string_character(ch) { + return Err(JsonRepairError::InvalidCharacter { + character: ch, + position: self.i, + }); + } + str_content.push(ch); + self.i += 1; + } + } + + if skip_escape_chars { + // repair: skip escape character + if self.current_char() == Some('\\') { + self.i += 1; + } + } + } + } + + fn parse_concatenated_string(&mut self) -> Result { + let mut processed = false; + + self.parse_whitespace_and_skip_comments(true); + while self.current_char() == Some('+') { + processed = true; + self.i += 1; + self.parse_whitespace_and_skip_comments(true); + + // repair: remove the end quote of the first string + self.output = self.strip_last_occurrence_with_remaining("\"", true); + let start = self.output.len(); + let parsed_str = self.parse_string(false, None)?; + if parsed_str { + // repair: remove the start quote of the second string + self.output = self.remove_at_index(start, 1); + } else { + // repair: remove the + because it's not followed by a string + self.output = self.insert_before_last_whitespace("\""); + } + } + + Ok(processed) + } + + fn parse_number(&mut self) -> Result { + let start = self.i; + + if self.current_char() == Some('-') { + self.i += 1; + if self.at_end_of_number() { + self.repair_number_ending_with_numeric_symbol(start); + return Ok(true); + } + if !self.is_digit(self.current_char()) { + self.i = start; + return Ok(false); + } + } + + while self.is_digit(self.current_char()) { + self.i += 1; + } + + if self.current_char() == Some('.') { + self.i += 1; + if self.at_end_of_number() { + self.repair_number_ending_with_numeric_symbol(start); + return Ok(true); + } + if !self.is_digit(self.current_char()) { + self.i = start; + return Ok(false); + } + while self.is_digit(self.current_char()) { + self.i += 1; + } + } + + if matches!(self.current_char(), Some('e' | 'E')) { + self.i += 1; + if matches!(self.current_char(), Some('-' | '+')) { + self.i += 1; + } + if self.at_end_of_number() { + self.repair_number_ending_with_numeric_symbol(start); + return Ok(true); + } + if !self.is_digit(self.current_char()) { + self.i = start; + return Ok(false); + } + while self.is_digit(self.current_char()) { + self.i += 1; + } + } + + if !self.at_end_of_number() { + self.i = start; + return Ok(false); + } + + if self.i > start { + let num: String = self + .chars + .get(start..self.i) + .map(|s| s.iter().collect()) + .unwrap_or_default(); + let has_invalid_leading_zero = + num.len() > 1 && num.starts_with('0') && self.is_digit(num.chars().nth(1)); + + if has_invalid_leading_zero { + self.output.push_str(&serde_json::to_string(&num).unwrap()); + } else { + self.output.push_str(&num); + } + return Ok(true); + } + + Ok(false) + } + + fn parse_keywords(&mut self) -> bool { + self.parse_keyword("true", "true") + || self.parse_keyword("false", "false") + || self.parse_keyword("null", "null") + || self.parse_keyword("True", "true") + || self.parse_keyword("False", "false") + || self.parse_keyword("None", "null") + } + + fn parse_keyword(&mut self, name: &str, value: &str) -> bool { + let name_chars: Vec = name.chars().collect(); + if self.i + name_chars.len() <= self.chars.len() { + let slice = self.chars.get(self.i..self.i + name_chars.len()); + if slice == Some(name_chars.as_slice()) { + self.output.push_str(value); + self.i += name_chars.len(); + return true; + } + } + false + } + + fn parse_unquoted_string(&mut self, is_key: bool) -> Result { + let start = self.i; + + if self.is_function_name_char_start(self.current_char()) { + while self.i < self.chars.len() && self.is_function_name_char(self.current_char()) { + self.i += 1; + } + + let mut j = self.i; + while j < self.chars.len() { + let ch = self.chars.get(j).copied().unwrap_or('\0'); + if !self.is_whitespace(ch, true) { + break; + } + j += 1; + } + + if j < self.chars.len() { + let ch_at_j = self.chars.get(j).copied().unwrap_or('\0'); + if ch_at_j == '(' { + // function call + self.i = j + 1; + self.parse_value()?; + if self.current_char() == Some(')') { + self.i += 1; + if self.current_char() == Some(';') { + self.i += 1; + } + } + return Ok(true); + } + } + } + + // reset for URL/unquoted string parsing + self.i = start; + + while self.i < self.chars.len() { + let ch = self.current_char().unwrap(); + if self.is_unquoted_string_delimiter(ch) + || self.is_quote(Some(ch)) + || (is_key && ch == ':') + { + break; + } + self.i += 1; + } + + // test for URL + if self.i > 0 + && self.chars.get(self.i - 1) == Some(&':') + && start < self.chars.len() + && self.i + 2 <= self.chars.len() + && self.is_url_start(self.chars.get(start..self.i + 2).unwrap_or(&[])) + { + while self.i < self.chars.len() && self.is_url_char(self.current_char().unwrap()) { + self.i += 1; + } + } + + if self.i > start { + // remove trailing whitespace + while self.i > start + && self.is_whitespace(self.chars.get(self.i - 1).copied().unwrap_or('\0'), true) + { + self.i -= 1; + } + + let symbol: String = self + .chars + .get(start..self.i) + .map(|s| s.iter().collect()) + .unwrap_or_default(); + + if symbol == "undefined" { + self.output.push_str("null"); + } else if is_key { + // Keys must always be strings in JSON + self.output + .push_str(&serde_json::to_string(&symbol).unwrap()); + } else { + // Try to parse as number first, if that fails, treat as string + if symbol.parse::().is_ok() { + // It's a valid number, output as-is + self.output.push_str(&symbol); + } else { + // It's a string, so quote it using serde_json + self.output + .push_str(&serde_json::to_string(&symbol).unwrap()); + } + } + + if self.current_char() == Some('"') { + // missing start quote, skip end quote + self.i += 1; + } + + return Ok(true); + } + + Ok(false) + } + + fn parse_regex(&mut self) -> Result { + if self.current_char() == Some('/') { + let start = self.i; + self.i += 1; + + while self.i < self.chars.len() + && (self.current_char() != Some('/') + || (self.i > 0 && self.chars.get(self.i - 1) == Some(&'\\'))) + { + self.i += 1; + } + + // Only increment if we found the closing slash + if self.i < self.chars.len() { + self.i += 1; + } + + let regex: String = self + .chars + .get(start..self.i) + .map(|s| s.iter().collect()) + .unwrap_or_default(); + self.output + .push_str(&serde_json::to_string(®ex).unwrap()); + return Ok(true); + } + Ok(false) + } + + // Helper methods + fn is_whitespace(&self, ch: char, skip_newline: bool) -> bool { + if skip_newline { + ch == ' ' || ch == '\n' || ch == '\t' || ch == '\r' + } else { + ch == ' ' || ch == '\t' || ch == '\r' + } + } + + fn is_special_whitespace(&self, ch: char) -> bool { + matches!(ch as u32, 0xa0 | 0x2000..=0x200a | 0x202f | 0x205f | 0x3000) + } + + fn is_quote(&self, ch: Option) -> bool { + matches!( + ch, + Some('"' | '\'' | '`' | '\u{2018}' | '\u{2019}' | '\u{201c}' | '\u{201d}' | '\u{00b4}') + ) + } + + fn is_digit(&self, ch: Option) -> bool { + matches!(ch, Some('0'..='9')) + } + + fn is_hex(&self, ch: char) -> bool { + ch.is_ascii_hexdigit() + } + + fn is_function_name_char_start(&self, ch: Option) -> bool { + matches!(ch, Some('a'..='z' | 'A'..='Z' | '_' | '$')) + } + + fn is_function_name_char(&self, ch: Option) -> bool { + matches!(ch, Some('a'..='z' | 'A'..='Z' | '0'..='9' | '_' | '$')) + } + + fn is_delimiter(&self, ch: Option) -> bool { + matches!( + ch, + Some(',' | ':' | '[' | ']' | '/' | '{' | '}' | '(' | ')' | '\n' | '+') + ) + } + + fn is_unquoted_string_delimiter(&self, ch: char) -> bool { + matches!(ch, ',' | '[' | ']' | '/' | '{' | '}' | '\n' | '+') + } + + fn is_start_of_value(&self, ch: Option) -> bool { + self.is_quote(ch) + || matches!( + ch, + Some('[' | '{' | 'a'..='z' | 'A'..='Z' | '0'..='9' | '_' | '-') + ) + } + + fn is_control_character(&self, ch: char) -> bool { + matches!(ch, '\n' | '\r' | '\t' | '\u{08}' | '\u{0c}') + } + + fn is_valid_string_character(&self, ch: char) -> bool { + ch >= '\u{0020}' + } + + fn is_url_start(&self, chars: &[char]) -> bool { + if chars.len() < 7 { + return false; + } + let s: String = chars.iter().collect(); + s.starts_with("http://") + || s.starts_with("https://") + || s.starts_with("ftp://") + || s.starts_with("mailto:") + || s.starts_with("file://") + || s.starts_with("data:") + || s.starts_with("irc://") + } + + fn is_url_char(&self, ch: char) -> bool { + ch.is_ascii_alphanumeric() + || matches!( + ch, + '-' | '.' + | '_' + | '~' + | ':' + | '/' + | '?' + | '#' + | '@' + | '!' + | '$' + | '&' + | '\'' + | '(' + | ')' + | '*' + | '+' + | ';' + | '=' + ) + } + + fn get_end_quote_matcher(&self, quote_char: char) -> fn(char) -> bool { + match quote_char { + '"' => |ch| ch == '"', + '\'' => |ch| ch == '\'', + '\u{2018}' => |ch| matches!(ch, '\u{2018}' | '\u{2019}'), + '\u{2019}' => |ch| matches!(ch, '\u{2018}' | '\u{2019}'), + '\u{201c}' => |ch| matches!(ch, '\u{201c}' | '\u{201d}'), + '\u{201d}' => |ch| matches!(ch, '\u{201c}' | '\u{201d}'), + '`' => |ch| matches!(ch, '`' | '\u{00b4}' | '\''), + '\u{00b4}' => |ch| matches!(ch, '`' | '\u{00b4}' | '\''), + _ => |ch| ch == '"', + } + } + + fn get_control_character_escape(&self, ch: char) -> String { + match ch { + '\u{08}' => "\\b".to_string(), + '\u{0c}' => "\\f".to_string(), + '\n' => "\\n".to_string(), + '\r' => "\\r".to_string(), + '\t' => "\\t".to_string(), + _ => ch.to_string(), + } + } + + fn at_end_of_block_comment(&self) -> bool { + self.current_char() == Some('*') && self.chars.get(self.i + 1) == Some(&'/') + } + + fn prev_non_whitespace_index(&self, start: usize) -> usize { + let mut prev = start; + while prev > 0 + && prev < self.chars.len() + && self.is_whitespace(self.chars.get(prev).copied().unwrap_or('\0'), true) + { + prev = prev.saturating_sub(1); + } + prev + } + + fn at_end_of_number(&self) -> bool { + self.i >= self.chars.len() + || self.is_delimiter(self.current_char()) + || self.is_whitespace(self.current_char().unwrap_or('\0'), true) + } + + fn repair_number_ending_with_numeric_symbol(&mut self, start: usize) { + let num: String = self + .chars + .get(start..self.i) + .map(|s| s.iter().collect()) + .unwrap_or_default(); + self.output.push_str(&format!("{num}0")); + } + + fn ends_with_comma_or_newline(&self) -> bool { + let regex = regex::Regex::new(r"[,\n][ \t\r]*$").unwrap(); + regex.is_match(&self.output) + } + + fn strip_last_occurrence(&self, text_to_strip: &str) -> String { + self.strip_last_occurrence_with_remaining(text_to_strip, false) + } + + fn strip_last_occurrence_with_remaining( + &self, + text_to_strip: &str, + strip_remaining: bool, + ) -> String { + if let Some(index) = self.output.rfind(text_to_strip) { + let mut result = self.output.get(..index).unwrap_or("").to_string(); + if !strip_remaining { + result.push_str(self.output.get(index + text_to_strip.len()..).unwrap_or("")); + } + result + } else { + self.output.clone() + } + } + + fn insert_before_last_whitespace(&self, text_to_insert: &str) -> String { + let chars: Vec = self.output.chars().collect(); + let mut index = chars.len(); + + if index == 0 || !self.is_whitespace(chars.get(index - 1).copied().unwrap_or('\0'), true) { + return format!("{}{}", self.output, text_to_insert); + } + + while index > 0 && self.is_whitespace(chars.get(index - 1).copied().unwrap_or('\0'), true) { + index -= 1; + } + + // Convert the char-based index back to a byte offset for string slicing. + let byte_index = self + .output + .char_indices() + .nth(index) + .map_or(self.output.len(), |(i, _)| i); + + format!( + "{}{}{}", + self.output.get(..byte_index).unwrap_or(""), + text_to_insert, + self.output.get(byte_index..).unwrap_or("") + ) + } + + fn insert_before_last_whitespace_str(&self, text: &str, text_to_insert: &str) -> String { + let chars: Vec = text.chars().collect(); + let mut index = chars.len(); + + if index == 0 || !self.is_whitespace(chars.get(index - 1).copied().unwrap_or('\0'), true) { + return format!("{text}{text_to_insert}"); + } + + while index > 0 && self.is_whitespace(chars.get(index - 1).copied().unwrap_or('\0'), true) { + index -= 1; + } + + // Convert the char-based index back to a byte offset for string slicing. + let byte_index = text + .char_indices() + .nth(index) + .map_or(text.len(), |(i, _)| i); + + format!( + "{}{}{}", + text.get(..byte_index).unwrap_or(""), + text_to_insert, + text.get(byte_index..).unwrap_or("") + ) + } + + fn remove_at_index(&self, start: usize, count: usize) -> String { + let chars: Vec = self.output.chars().collect(); + let mut result = String::new(); + for (i, ch) in chars.iter().enumerate() { + if i < start || i >= start + count { + result.push(*ch); + } + } + result + } +} + +pub fn json_repair Deserialize<'de>>(text: &str) -> Result { + let parser = JsonRepairParser::new(text.to_string()); + parser.parse() +} diff --git a/crates/forge_json_repair/src/schema_coercion.rs b/crates/forge_json_repair/src/schema_coercion.rs new file mode 100644 index 0000000000000000000000000000000000000000..48fc0901f818158410fadceee850ceb33f31503a --- /dev/null +++ b/crates/forge_json_repair/src/schema_coercion.rs @@ -0,0 +1,1397 @@ +use schemars::Schema; +use serde::de::Error as _; +use serde_json::Value; + +use crate::json_repair; + +/// Coerces a JSON value to match the expected types defined in a JSON schema. +/// +/// This function recursively traverses the JSON value and the schema, +/// converting string values to the expected types (e.g., "42" -> 42) when the +/// schema indicates a different type is expected. +/// +/// # Arguments +/// +/// * `value` - The JSON value to coerce +/// * `schema` - The JSON schema defining expected types +/// +/// # Errors +/// +/// Returns the original value if coercion is not possible or the schema doesn't +/// specify type constraints. +pub fn coerce_to_schema(value: Value, schema: &Schema) -> Value { + coerce_value_with_schema(value, schema, schema) +} + +fn coerce_value_with_schema(value: Value, schema: &Schema, root_schema: &Schema) -> Value { + // In schemars 1.0, Schema is a wrapper around serde_json::Value + // It can be either a boolean or an object + match schema.as_value() { + Value::Bool(_) => value, // Boolean schemas don't provide type info for coercion + Value::Object(schema_obj) => { + coerce_value_with_schema_object(value, schema_obj, root_schema) + } + _ => value, // Shouldn't happen, but handle gracefully + } +} + +fn coerce_value_with_schema_object( + value: Value, + schema: &serde_json::Map, + root_schema: &Schema, +) -> Value { + // Handle $ref schemas by resolving references + if let Some(Value::String(reference)) = schema.get("$ref") { + // Resolve $ref against root schema definitions + // schemars uses format: "#/$defs/TypeName" or "#/definitions/TypeName" + if let Some(def_name) = reference + .strip_prefix("#/$defs/") + .or_else(|| reference.strip_prefix("#/definitions/")) + && let Some(def_schema) = root_schema + .as_value() + .as_object() + .and_then(|obj| obj.get("$defs").or_else(|| obj.get("definitions"))) + .and_then(|defs| defs.as_object()) + .and_then(|defs| defs.get(def_name)) + { + // Convert the Value to a Schema wrapper + if let Ok(schema) = serde_json::from_value::(def_schema.clone()) { + return coerce_value_with_schema(value, &schema, root_schema); + } + } + } + + // Coerce empty strings to null for nullable schemas. + // LLMs often send "" for optional parameters instead of omitting them or + // sending null. When the schema has "nullable": true (OpenAPI 3.0 style), + // an empty string should be treated as null. + if let Value::String(s) = &value + && s.is_empty() + && is_nullable(schema) + { + return Value::Null; + } + // Handle anyOf/oneOf schemas by trying each sub-schema + if let Some(any_of) = schema.get("anyOf").and_then(|v| v.as_array()) { + // Try each sub-schema in anyOf until one succeeds + for sub_schema_value in any_of { + if let Ok(sub_schema) = serde_json::from_value::(sub_schema_value.clone()) { + let result = coerce_value_with_schema(value.clone(), &sub_schema, root_schema); + if result != value { + return result; + } + } + } + } + if let Some(one_of) = schema.get("oneOf").and_then(|v| v.as_array()) { + // Try each sub-schema in oneOf until one succeeds + for sub_schema_value in one_of { + if let Ok(sub_schema) = serde_json::from_value::(sub_schema_value.clone()) { + let result = coerce_value_with_schema(value.clone(), &sub_schema, root_schema); + if result != value { + return result; + } + } + } + } + if let Some(all_of) = schema.get("allOf").and_then(|v| v.as_array()) { + // Apply all schemas in sequence + let mut result = value; + for sub_schema_value in all_of { + if let Ok(sub_schema) = serde_json::from_value::(sub_schema_value.clone()) { + result = coerce_value_with_schema(result, &sub_schema, root_schema); + } + } + return result; + } + + // Handle objects with properties + if let Value::Object(mut map) = value { + if let Some(Value::Object(properties)) = schema.get("properties") { + for (key, val) in map.iter_mut() { + if let Some(prop_schema_value) = properties.get(key) + && let Ok(prop_schema) = + serde_json::from_value::(prop_schema_value.clone()) + { + let coerced = coerce_value_with_schema(val.clone(), &prop_schema, root_schema); + *val = coerced; + } + } + } + return Value::Object(map); + } + + // Handle arrays + if let Value::Array(arr) = value { + // Check for prefixItems first (JSON Schema 2020-12 for tuples) + if let Some(prefix_items_value) = schema.get("prefixItems") + && let Some(item_schemas) = prefix_items_value.as_array() + { + // Array of schemas (tuple validation) + return Value::Array( + arr.into_iter() + .enumerate() + .map(|(i, item)| { + if let Some(schema_value) = item_schemas.get(i) + && let Ok(schema) = + serde_json::from_value::(schema_value.clone()) + { + return coerce_value_with_schema(item, &schema, root_schema); + } + item + }) + .collect(), + ); + } + + // Check for items (older JSON Schema drafts) + if let Some(items_schema_value) = schema.get("items") { + // Check if it's a single schema or array of schemas + if items_schema_value.is_object() || items_schema_value.is_boolean() { + // Single schema for all items + if let Ok(item_schema) = + serde_json::from_value::(items_schema_value.clone()) + { + return Value::Array( + arr.into_iter() + .map(|item| coerce_value_with_schema(item, &item_schema, root_schema)) + .collect(), + ); + } + } else if let Some(item_schemas) = items_schema_value.as_array() { + // Array of schemas (tuple validation) + return Value::Array( + arr.into_iter() + .enumerate() + .map(|(i, item)| { + if let Some(schema_value) = item_schemas.get(i) + && let Ok(schema) = + serde_json::from_value::(schema_value.clone()) + { + return coerce_value_with_schema(item, &schema, root_schema); + } + item + }) + .collect(), + ); + } + } + return Value::Array(arr); + } + + // If schema has specific instance types, try to coerce the value + if let Some(type_value) = schema.get("type") { + return coerce_by_type(value, type_value, schema, root_schema); + } + + value +} + +fn coerce_by_type( + value: Value, + type_value: &Value, + schema: &serde_json::Map, + root_schema: &Schema, +) -> Value { + // type can be a string or an array of strings + let target_types: Vec<&str> = match type_value { + Value::String(s) => vec![s.as_str()], + Value::Array(arr) => arr.iter().filter_map(|v| v.as_str()).collect(), + _ => return value, + }; + + // If the value already matches one of the target types, return as-is + if type_matches(&value, &target_types) { + return value; + } + + // Try coercion if value is a string + if let Value::String(s) = &value { + for target_type in target_types { + if let Some(coerced) = try_coerce_string(s, target_type, schema, root_schema) { + return coerced; + } + } + } + + value +} + +/// Checks if a schema is marked as nullable via the OpenAPI 3.0 "nullable" +/// extension. This is set by schemars when `option_nullable = true` for +/// `Option` fields. +fn is_nullable(schema: &serde_json::Map) -> bool { + schema + .get("nullable") + .and_then(|v| v.as_bool()) + .unwrap_or(false) +} + +fn type_matches(value: &Value, target_types: &[&str]) -> bool { + target_types.iter().any(|t| match *t { + "null" => value.is_null(), + "boolean" => value.is_boolean(), + "object" => value.is_object(), + "array" => value.is_array(), + "number" => value.is_number(), + "string" => value.is_string(), + "integer" => value.is_i64() || value.is_u64(), + _ => false, + }) +} + +fn try_coerce_string( + s: &str, + target_type: &str, + schema: &serde_json::Map, + root_schema: &Schema, +) -> Option { + match target_type { + "integer" => { + // Try to parse as i64 + if let Ok(num) = s.parse::() { + return Some(Value::Number(num.into())); + } + // Try to parse as u64 + if let Ok(num) = s.parse::() { + return Some(Value::Number(num.into())); + } + None + } + "number" => { + // Try to parse as integer first + if let Ok(num) = s.parse::() { + return Some(Value::Number(num.into())); + } + // Then try float + if let Ok(num) = s.parse::() + && let Some(json_num) = serde_json::Number::from_f64(num) + { + return Some(Value::Number(json_num)); + } + None + } + "boolean" => match s.trim().to_lowercase().as_str() { + "true" => Some(Value::Bool(true)), + "false" => Some(Value::Bool(false)), + _ => None, + }, + "null" => { + if s.trim().to_lowercase() == "null" { + Some(Value::Null) + } else { + None + } + } + "string" => { + // Keep as string + None + } + "object" => { + // Try to parse the string as a JSON object + if let Ok(parsed) = try_parse_json_string(s) + && parsed.is_object() + { + return Some(parsed); + } + None + } + "array" => { + // Try to parse the string as a JSON array + if let Ok(parsed) = try_parse_json_string(s) + && parsed.is_array() + { + // Recursively coerce array items using the schema + return Some(coerce_array_value(parsed, schema, root_schema)); + } + + // If direct parsing fails, try to extract array portion from the string + // This handles cases like: "[\"item\"]{\n}" or "garbage[\"item\"]more" + if let Some(extracted) = extract_array_from_string(s) { + // Recursively coerce the extracted array items + return Some(coerce_array_value(extracted, schema, root_schema)); + } + + None + } + _ => None, + } +} + +/// Helper function to recursively coerce array items based on the schema +fn coerce_array_value( + value: Value, + schema: &serde_json::Map, + root_schema: &Schema, +) -> Value { + if let Value::Array(arr) = value { + // Check for prefixItems first (JSON Schema 2020-12 for tuples) + if let Some(prefix_items_value) = schema.get("prefixItems") + && let Some(item_schemas) = prefix_items_value.as_array() + { + // Array of schemas (tuple validation) + return Value::Array( + arr.into_iter() + .enumerate() + .map(|(i, item)| { + if let Some(schema_value) = item_schemas.get(i) + && let Ok(schema) = + serde_json::from_value::(schema_value.clone()) + { + return coerce_value_with_schema(item, &schema, root_schema); + } + item + }) + .collect(), + ); + } + + // Check if schema defines array item types (older JSON Schema draft) + if let Some(items_schema_value) = schema.get("items") { + // Check if it's a single schema or array of schemas + if items_schema_value.is_object() || items_schema_value.is_boolean() { + // Single schema for all items + if let Ok(item_schema) = + serde_json::from_value::(items_schema_value.clone()) + { + return Value::Array( + arr.into_iter() + .map(|item| coerce_value_with_schema(item, &item_schema, root_schema)) + .collect(), + ); + } + } else if let Some(item_schemas) = items_schema_value.as_array() { + // Array of schemas (tuple validation) + return Value::Array( + arr.into_iter() + .enumerate() + .map(|(i, item)| { + if let Some(schema_value) = item_schemas.get(i) + && let Ok(schema) = + serde_json::from_value::(schema_value.clone()) + { + return coerce_value_with_schema(item, &schema, root_schema); + } + item + }) + .collect(), + ); + } + } + Value::Array(arr) + } else { + value + } +} + +/// Attempts to parse a string as JSON, JSON5, or repairable JSON, and unwraps +/// nested JSON strings when needed. +fn try_parse_json_string(s: &str) -> Result { + let mut parsed = parse_json_like_value(s)?; + + for _ in 0..4 { + let Value::String(inner) = &parsed else { + return Ok(parsed); + }; + + let Ok(next) = parse_json_like_value(inner) else { + return Ok(parsed); + }; + + parsed = next; + } + + Ok(parsed) +} + +fn parse_json_like_value(s: &str) -> Result { + // First try parsing as-is (valid JSON) + if let Ok(parsed) = serde_json::from_str::(s) { + return Ok(parsed); + } + + // If that fails, try parsing as JSON5 (handles single quotes, comments, etc.) + if let Ok(parsed) = serde_json5::from_str::(s) { + return Ok(parsed); + } + + // Finally, fall back to Forge's JSON repair for malformed-but-recoverable + // payloads such as persisted double-encoded tool arguments. + json_repair(s).map_err(|e| serde_json::Error::custom(e.to_string())) +} + +/// Extracts an array from a string that may contain garbage before/after the +/// array +/// +/// # Examples +/// +/// - `"[\"item\"]{\n}"` -> `["item"]` +/// - `"garbage[\"item\"]"` -> `["item"]` +/// - `"prefix[1,2,3]suffix"` -> `[1,2,3]` +/// +/// This function is more permissive than standard JSON parsing - it will +/// extract arrays that have trailing or leading garbage. It requires the string +/// to at least look like it contains array-like content (quotes, commas, or +/// brackets after '['). +fn extract_array_from_string(s: &str) -> Option { + // Find the first '[' and try to extract array from there + let start_idx = s.find('[')?; + + // Check if there's anything after '[' that looks like array content + // This helps us avoid extracting arrays from clearly invalid strings like + // "[invalid json" + let after_bracket = s.get(start_idx + 1..).unwrap_or(""); + let has_array_like_content = after_bracket.contains('"') + || after_bracket.contains(',') + || after_bracket.contains(']') + || after_bracket.chars().next().is_some_and(|c| c.is_numeric()); + + if !has_array_like_content { + return None; + } + + // Try to find matching closing bracket by parsing incrementally + // Start from the opening bracket and try increasingly shorter substrings. + // We iterate over valid char boundaries to avoid panicking on multi-byte + // UTF-8 characters where byte offsets can land inside a character. + for (end_idx, _) in s.char_indices().rev() { + if end_idx <= start_idx { + break; + } + let Some(candidate) = s.get(start_idx..end_idx) else { + continue; + }; + + // Try to repair and parse this candidate + if let Ok(parsed) = crate::json_repair::(candidate) + && parsed.is_array() + { + return Some(parsed); + } + } + + // Also try the full string as a last resort (end at s.len() which is + // always a valid boundary) + let candidate = s.get(start_idx..)?; + if let Ok(parsed) = crate::json_repair::(candidate) + && parsed.is_array() + { + return Some(parsed); + } + + None +} + +#[cfg(test)] +mod tests { + #![allow(dead_code)] + use schemars::{JsonSchema, schema_for}; + use serde::{Deserialize, Serialize}; + use serde_json::json; + + use super::*; + + // Test structs with JsonSchema derive + #[derive(JsonSchema)] + #[allow(dead_code)] + struct AgeData { + age: i64, + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct RangeData { + start: i64, + end: i64, + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct PriceData { + price: f64, + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct BooleanData { + active: bool, + disabled: bool, + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct UserData { + age: i64, + score: f64, + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct UserWrapper { + user: UserData, + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct NumbersData { + numbers: Vec, + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct MixedData { + name: String, + age: i64, + active: bool, + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct PathData { + path: String, + start_line: i64, + end_line: i64, + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct IntOrNull { + value: Option, + } + + #[derive(JsonSchema, Deserialize, Serialize)] + #[allow(dead_code)] + #[serde(untagged)] + enum IntOrBool { + Int(i64), + Bool(bool), + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct IntOrBoolData { + value: IntOrBool, + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct AllOfIntNumber { + value: i64, + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct CoordinatesData { + coordinates: [f64; 3], + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct MixedTupleData { + data: (String, i64, bool), + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct TupleItems { + items: [i64; 2], + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct ExtraItemsData { + items: Vec, + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct NestedUnionData { + nested: IntOrNull, + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct NullData { + value: Option<()>, + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct BoolData { + value: bool, + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct LargeIntData { + value: i64, + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct UnsignedIntData { + value: u64, + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct ArrayData { + items: Vec, + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct EditsData { + edits: Vec, + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct ConfigData { + config: std::collections::BTreeMap, + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct DataArray { + data: Vec, + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct ItemsArray { + items: Vec, + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct ConfigWithComments { + config: std::collections::BTreeMap, + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct ItemsTrailingComma { + items: Vec, + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct MultiPatchData { + edits: Vec, + } + + #[test] + fn test_coerce_string_to_integer() { + let fixture = json!({"age": "42"}); + let schema = schema_for!(AgeData); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"age": 42}); + assert_eq!(actual, expected); + } + + #[test] + fn test_coerce_multiple_string_integers() { + let fixture = json!({"start": "100", "end": "200"}); + let schema = schema_for!(RangeData); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"start": 100, "end": 200}); + assert_eq!(actual, expected); + } + + #[test] + fn test_coerce_string_to_number_float() { + let fixture = json!({"price": "19.99"}); + let schema = schema_for!(PriceData); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"price": 19.99}); + assert_eq!(actual, expected); + } + + #[test] + fn test_coerce_string_to_boolean() { + let fixture = json!({"active": "true", "disabled": "false"}); + let schema = schema_for!(BooleanData); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"active": true, "disabled": false}); + assert_eq!(actual, expected); + } + + #[test] + fn test_no_coercion_when_types_match() { + let fixture = json!({"age": 42}); + let schema = schema_for!(AgeData); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"age": 42}); + assert_eq!(actual, expected); + } + + #[test] + fn test_no_coercion_for_invalid_strings() { + let fixture = json!({"age": "not_a_number"}); + let schema = schema_for!(AgeData); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"age": "not_a_number"}); + assert_eq!(actual, expected); + } + + #[test] + fn test_coerce_nested_objects() { + let fixture = json!({"user": {"age": "30", "score": "95.5"}}); + let schema = schema_for!(UserWrapper); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"user": {"age": 30, "score": 95.5}}); + assert_eq!(actual, expected); + } + + #[test] + fn test_coerce_array_items() { + let fixture = json!({"numbers": ["1", "2", "3"]}); + let schema = schema_for!(NumbersData); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"numbers": [1, 2, 3]}); + assert_eq!(actual, expected); + } + + #[test] + fn test_preserve_non_string_values() { + let fixture = json!({"name": "John", "age": 42, "active": true}); + let schema = schema_for!(MixedData); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"name": "John", "age": 42, "active": true}); + assert_eq!(actual, expected); + } + + #[test] + fn test_read_tool_line_numbers() { + // Simulate the exact case from the task: read tool with string line numbers + let fixture = json!({ + "path": "/Users/amit/code-forge/crates/forge_main/src/ui.rs", + "start_line": "2255", + "end_line": "2285" + }); + + let schema = schema_for!(PathData); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({ + "path": "/Users/amit/code-forge/crates/forge_main/src/ui.rs", + "start_line": 2255, + "end_line": 2285 + }); + assert_eq!(actual, expected); + } + + #[test] + fn test_coerce_any_of_union_types() { + // Test coercing string to integer + let fixture = json!({"value": "42"}); + let schema = schema_for!(IntOrNull); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"value": 42}); + assert_eq!(actual, expected); + + // Test preserving null + let fixture = json!({"value": null}); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"value": null}); + assert_eq!(actual, expected); + } + + #[test] + fn test_coerce_one_of_union_types() { + // Test coercing string to integer + let fixture = json!({"value": "123"}); + let schema = schema_for!(IntOrBoolData); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"value": 123}); + assert_eq!(actual, expected); + + // Test coercing string to boolean + let fixture = json!({"value": "true"}); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"value": true}); + assert_eq!(actual, expected); + } + + #[test] + fn test_coerce_all_of_composition() { + // Test coercing string to integer via allOf composition + let fixture = json!({"value": "42"}); + let schema = schema_for!(AllOfIntNumber); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"value": 42}); + assert_eq!(actual, expected); + } + + #[test] + fn test_any_of_preserves_original_when_no_match() { + // Test that anyOf preserves original value when no subschema matches + // Note: oneOf behaves similarly + let fixture = json!({"value": "not_a_number"}); + let schema = schema_for!(IntOrBoolData); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"value": "not_a_number"}); + assert_eq!(actual, expected); + } + + #[test] + fn test_any_of_with_number_coercion() { + // Test anyOf with number coercion + let fixture = json!({"value": "2.14"}); + let schema = schema_for!(IntOrNull); + let actual = coerce_to_schema(fixture, &schema); + // The anyOf schema tries each subschema; since "2.14" can't be parsed as i64, + // it returns the original value + let expected = json!({"value": "2.14"}); + assert_eq!(actual, expected); + } + + #[test] + fn test_array_with_tuple_schema() { + // Test array with tuple schema (SingleOrVec::Vec) + let fixture = json!({"coordinates": ["1.5", "2.5", "3.5"]}); + let schema = schema_for!(CoordinatesData); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"coordinates": [1.5, 2.5, 3.5]}); + assert_eq!(actual, expected); + } + + #[test] + fn test_array_with_tuple_schema_mixed_types() { + // Test array with tuple schema with mixed types + let fixture = json!({"data": ["name", "42", "true"]}); + let schema = schema_for!(MixedTupleData); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"data": ["name", 42, true]}); + assert_eq!(actual, expected); + } + + #[test] + fn test_array_with_tuple_schema_extra_items() { + // Test that Vec doesn't coerce items (no type constraints) + let fixture = json!({"items": ["1", "2", "3", "4"]}); + let schema = schema_for!(ExtraItemsData); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"items": ["1", "2", "3", "4"]}); + assert_eq!(actual, expected); + } + + #[test] + fn test_nested_any_of_in_object() { + // Test coercing in nested object with anyOf + let fixture = json!({"nested": {"value": "42"}}); + let schema = schema_for!(NestedUnionData); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"nested": {"value": 42}}); + assert_eq!(actual, expected); + } + + #[test] + fn test_coerce_string_to_null() { + // Test coercing "null" string to null + let fixture = json!({"value": "null"}); + let schema = schema_for!(NullData); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"value": null}); + assert_eq!(actual, expected); + + // Test that "NULL" (uppercase) also works + let fixture = json!({"value": "NULL"}); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"value": null}); + assert_eq!(actual, expected); + } + + #[test] + fn test_coerce_boolean_case_insensitive() { + // Test that boolean coercion is case-insensitive + let schema = schema_for!(BoolData); + + // Test various case variations + for (input, expected) in [ + ("true", true), + ("TRUE", true), + ("True", true), + ("false", false), + ("FALSE", false), + ("False", false), + ] { + let fixture = json!({"value": input}); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"value": expected}); + assert_eq!(actual, expected); + } + } + + #[test] + fn test_coerce_large_integer() { + // Test coercing large integers that fit in i64 + let schema = schema_for!(LargeIntData); + + // Test coercing large positive integer + let fixture = json!({"value": "9223372036854775807"}); // i64::MAX + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"value": 9223372036854775807i64}); + assert_eq!(actual, expected); + + // Test coercing large negative integer + let fixture = json!({"value": "-9223372036854775808"}); // i64::MIN + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"value": -9223372036854775808i64}); + assert_eq!(actual, expected); + } + + #[test] + fn test_coerce_unsigned_integer() { + // Test coercing unsigned integers (u64) + let schema = schema_for!(UnsignedIntData); + + // Test coercing large unsigned integer that doesn't fit in i64 + let fixture = json!({"value": "18446744073709551615"}); // u64::MAX + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"value": 18446744073709551615u64}); + assert_eq!(actual, expected); + } + + #[test] + fn test_coerce_string_to_array() { + // Test coercing a JSON array string to an actual array + let fixture = json!({"items": "[1, 2, 3]"}); + let schema = schema_for!(ArrayData); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"items": [1, 2, 3]}); + assert_eq!(actual, expected); + } + + #[test] + fn test_coerce_python_style_string_to_array() { + // Test coercing a Python-style array string to an actual array + let fixture = json!({"edits": "[{'content': 'test', 'operation': 'replace'}]"}); + let schema = schema_for!(EditsData); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"edits": [{"content": "test", "operation": "replace"}]}); + assert_eq!(actual, expected); + } + + #[test] + fn test_coerce_python_style_string_to_object() { + // Test coercing a Python-style object string to an actual object + let fixture = json!({"config": "{'key': 'value', 'number': 42}"}); + let schema = schema_for!(ConfigData); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"config": {"key": "value", "number": 42}}); + assert_eq!(actual, expected); + } + + #[test] + fn test_coerce_double_encoded_string_to_object() { + let fixture = json!({"config": r#""{\"key\":\"value\",\"number\":42}""#}); + let schema = schema_for!(ConfigData); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"config": {"key": "value", "number": 42}}); + assert_eq!(actual, expected); + } + + #[test] + fn test_repairs_invalid_json_string_when_schema_expects_array() { + // Invalid JSON-like array strings are repaired into arrays when the schema + // expects one. + let fixture = json!({"data": "[invalid json"}); + let schema = schema_for!(DataArray); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"data": ["invalid json"]}); + assert_eq!(actual, expected); + } + + #[test] + fn test_coerce_json5_with_comments() { + // Test coercing JSON5 with comments + let fixture = json!({"config": r#"{ + // This is a comment + "key": "value", + "number": 42, + }"#}); + let schema = schema_for!(ConfigWithComments); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"config": {"key": "value", "number": 42}}); + assert_eq!(actual, expected); + } + + #[test] + fn test_coerce_json5_with_trailing_commas() { + // Test coercing JSON5 with trailing commas + let fixture = json!({"items": "[1, 2, 3,]"}); + let schema = schema_for!(ItemsTrailingComma); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"items": [1, 2, 3]}); + assert_eq!(actual, expected); + } + + #[test] + fn test_coerce_multi_patch_python_style() { + // Test coercing exact Python-style input from error + // This matches multi_patch tool call format with nested objects + let python_style = r#"[{'content': 'use schemars::schema::{InstanceType, RootSchema, Schema, SchemaObject, SingleOrVec};', 'operation': 'replace', 'path': 'crates/forge_json_repair/src/schema_coercion.rs'}, {'content': 'fn coerce_value_with_schema(value: Value, schema: &Schema) -> Value {', 'operation': 'replace', 'path': 'crates/forge_json_repair/src/schema_coercion.rs'}]"#; + + let fixture = json!({"edits": python_style}); + let schema = schema_for!(MultiPatchData); + let actual = coerce_to_schema(fixture, &schema); + + // Should coerce string to an array of objects + assert!(actual["edits"].is_array()); + let edits = actual["edits"].as_array().unwrap(); + assert_eq!(edits.len(), 2); + + // Verify first edit object + assert_eq!( + edits[0]["content"], + "use schemars::schema::{InstanceType, RootSchema, Schema, SchemaObject, SingleOrVec};" + ); + assert_eq!(edits[0]["operation"], "replace"); + assert_eq!( + edits[0]["path"], + "crates/forge_json_repair/src/schema_coercion.rs" + ); + + // Verify second edit object + assert_eq!( + edits[1]["content"], + "fn coerce_value_with_schema(value: Value, schema: &Schema) -> Value {" + ); + assert_eq!(edits[1]["operation"], "replace"); + assert_eq!( + edits[1]["path"], + "crates/forge_json_repair/src/schema_coercion.rs" + ); + } + + // Tests for array extraction with garbage + #[derive(JsonSchema)] + #[allow(dead_code)] + struct AgentInput { + tasks: Vec, + } + + #[test] + fn test_coerce_malformed_string_array_with_trailing_garbage() { + // This is the exact case from the issue: string that looks like an array but + // has trailing garbage + let fixture = json!({ + "tasks": "[\"Find where the main function is defined in the code-forge codebase. Search for main function definitions and entry points.\"]{\n}" + }); + + let schema = schema_for!(AgentInput); + let actual = coerce_to_schema(fixture, &schema); + + // Should extract the array portion and ignore the trailing garbage + let expected = json!({ + "tasks": ["Find where the main function is defined in the code-forge codebase. Search for main function definitions and entry points."] + }); + + assert_eq!(actual, expected); + } + + #[test] + fn test_coerce_string_array_with_leading_garbage() { + // Test with leading garbage before the array + let fixture = json!({ + "tasks": "garbage[\"task1\", \"task2\"]" + }); + + let schema = schema_for!(AgentInput); + let actual = coerce_to_schema(fixture, &schema); + + // Should extract the array portion from the string + let expected = json!({ + "tasks": ["task1", "task2"] + }); + + assert_eq!(actual, expected); + } + + #[test] + fn test_coerce_string_array_with_both_garbage() { + // Test with garbage on both ends + let fixture = json!({ + "tasks": "prefix[\"task1\"]suffix" + }); + + let schema = schema_for!(AgentInput); + let actual = coerce_to_schema(fixture, &schema); + + // Should extract the array portion + let expected = json!({ + "tasks": ["task1"] + }); + + assert_eq!(actual, expected); + } + + #[test] + fn test_repair_then_coerce_full_payload() { + // Test the full flow: repair the JSON structure, then coerce types + let malformed = r#"{"tasks": "[\"Find main function\"]{\n}"}"#; + + // First repair the JSON structure + let repaired: Value = crate::json_repair(malformed).expect("Should repair JSON"); + + // Then coerce to schema + let schema = schema_for!(AgentInput); + let actual = coerce_to_schema(repaired, &schema); + + let expected = json!({ + "tasks": ["Find main function"] + }); + + assert_eq!(actual, expected); + } + + #[test] + fn test_preserve_already_valid_array() { + // Ensure we don't break valid arrays + let fixture = json!({ + "tasks": ["task1", "task2"] + }); + + let schema = schema_for!(AgentInput); + let actual = coerce_to_schema(fixture, &schema); + + let expected = json!({ + "tasks": ["task1", "task2"] + }); + + assert_eq!(actual, expected); + } + + #[test] + fn test_coerce_string_array_without_garbage() { + // Valid JSON array string + let fixture = json!({ + "tasks": "[\"task1\", \"task2\"]" + }); + + let schema = schema_for!(AgentInput); + let actual = coerce_to_schema(fixture, &schema); + + let expected = json!({ + "tasks": ["task1", "task2"] + }); + + assert_eq!(actual, expected); + } + + // Test nested structures + #[derive(JsonSchema)] + #[allow(dead_code)] + struct SearchQuery { + query: String, + use_case: String, + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct SemanticSearchInput { + queries: Vec, + extensions: Vec, + } + + #[test] + fn test_coerce_nested_array_of_objects_with_garbage() { + // Test array of objects with trailing garbage + let fixture = json!({ + "queries": "[{\"query\": \"test\", \"use_case\": \"find\"}]garbage", + "extensions": "[\".rs\"]" + }); + + let schema = schema_for!(SemanticSearchInput); + let actual = coerce_to_schema(fixture, &schema); + + let expected = json!({ + "queries": [{"query": "test", "use_case": "find"}], + "extensions": [".rs"] + }); + + assert_eq!(actual, expected); + } + + #[test] + fn test_coerce_nested_array_with_string_numbers() { + // Test that nested coercion works - string numbers inside objects inside arrays + #[derive(JsonSchema)] + #[allow(dead_code)] + struct Item { + id: i64, + name: String, + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct ItemList { + items: Vec, + } + + let fixture = json!({ + "items": "[{\"id\": \"42\", \"name\": \"test\"}]extra" + }); + + let schema = schema_for!(ItemList); + let actual = coerce_to_schema(fixture, &schema); + + // The id should be coerced from string "42" to number 42 + let expected = json!({ + "items": [{"id": 42, "name": "test"}] + }); + + assert_eq!(actual, expected); + } + + #[test] + fn test_coerce_deeply_nested_arrays() { + // Test arrays containing objects containing arrays + #[derive(JsonSchema)] + #[allow(dead_code)] + struct DeepItem { + tags: Vec, + } + + #[derive(JsonSchema)] + #[allow(dead_code)] + struct DeepList { + items: Vec, + } + + let fixture = json!({ + "items": "[{\"tags\": [\"tag1\", \"tag2\"]}]garbage" + }); + + let schema = schema_for!(DeepList); + let actual = coerce_to_schema(fixture, &schema); + + let expected = json!({ + "items": [{"tags": ["tag1", "tag2"]}] + }); + + assert_eq!(actual, expected); + } + + #[test] + fn test_coerce_empty_string_to_null_for_nullable_field() { + // Simulates LLM sending "" for a nullable string field (e.g., file_type in + // fs_search). The schema uses "nullable: true" (OpenAPI 3.0 style). + #[derive(JsonSchema)] + #[allow(dead_code)] + struct NullableStringData { + required_field: String, + #[schemars(default)] + optional_field: Option, + } + + // Generate schema with option_nullable=true (matching project settings) + let settings = schemars::generate::SchemaSettings::default() + .with_transform(schemars::transform::AddNullable::default()) + .into_generator(); + let schema = settings.into_root_schema_for::(); + + let fixture = json!({ + "required_field": "value", + "optional_field": "" + }); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({ + "required_field": "value", + "optional_field": null + }); + assert_eq!(actual, expected); + } + + #[test] + fn test_preserve_non_empty_string_for_nullable_field() { + // Non-empty strings should NOT be converted to null, even for nullable fields + #[derive(JsonSchema)] + #[allow(dead_code)] + struct NullableStringData { + optional_field: Option, + } + + let settings = schemars::generate::SchemaSettings::default() + .with_transform(schemars::transform::AddNullable::default()) + .into_generator(); + let schema = settings.into_root_schema_for::(); + let fixture = json!({"optional_field": "rust"}); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"optional_field": "rust"}); + assert_eq!(actual, expected); + } + + #[test] + fn test_preserve_empty_string_for_required_field() { + // Empty strings should NOT be converted to null for non-nullable fields + #[derive(JsonSchema)] + #[allow(dead_code)] + struct RequiredStringData { + name: String, + } + + let schema = schema_for!(RequiredStringData); + + let fixture = json!({"name": ""}); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"name": ""}); + assert_eq!(actual, expected); + } + + #[test] + fn test_coerce_empty_string_to_null_for_nullable_integer() { + // Empty string for a nullable integer should become null + #[derive(JsonSchema)] + #[allow(dead_code)] + struct NullableIntData { + count: Option, + } + + let settings = schemars::generate::SchemaSettings::default() + .with_transform(schemars::transform::AddNullable::default()) + .into_generator(); + let schema = settings.into_root_schema_for::(); + + let fixture = json!({"count": ""}); + let actual = coerce_to_schema(fixture, &schema); + let expected = json!({"count": null}); + assert_eq!(actual, expected); + } + + #[test] + fn test_extract_array_from_string_with_multibyte_chars() { + // Multi-byte UTF-8 characters (like arrows and emojis) should not + // cause panics when extract_array_from_string iterates over byte + // positions. The function must only slice at valid char boundaries. + let input = "prefix → [1, 2, 3] suffix"; + let result = extract_array_from_string(input); + assert!(result.is_some()); + let arr = result.unwrap(); + assert!(arr.is_array()); + assert_eq!(arr.as_array().unwrap().len(), 3); + } + + #[test] + fn test_extract_array_from_string_with_emoji_prefix() { + // Emoji characters are 4 bytes each, many byte positions inside them + // are invalid char boundaries. + let input = "🔑🔒 [4, 5, 6]"; + let result = extract_array_from_string(input); + assert!(result.is_some()); + let arr = result.unwrap(); + assert!(arr.is_array()); + assert_eq!(arr.as_array().unwrap().len(), 3); + } + + #[test] + fn test_extract_array_from_string_with_multibyte_inside_array() { + // Multi-byte chars inside the array value itself + let input = r#"["αβγ", "δεζ"]"#; + let result = extract_array_from_string(input); + assert!(result.is_some()); + let arr = result.unwrap(); + assert!(arr.is_array()); + } +} diff --git a/crates/forge_json_repair/tests/comments_unicode.rs b/crates/forge_json_repair/tests/comments_unicode.rs new file mode 100644 index 0000000000000000000000000000000000000000..5b5a9aaaee9be6866913b4f3cb085a99b20b5b20 --- /dev/null +++ b/crates/forge_json_repair/tests/comments_unicode.rs @@ -0,0 +1,57 @@ +use forge_json_repair::json_repair; +use pretty_assertions::assert_eq; + +#[test] +fn test_remove_comments() { + let fixture = r#"/* foo */ {}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({}); + assert_eq!(actual, expected); + + let fixture = r#"{} /* foo */"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({}); + assert_eq!(actual, expected); + + let fixture = r#"{"a":"foo",/*hello*/"b":"bar"}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": "foo", "b": "bar"}); + assert_eq!(actual, expected); + + let fixture = r#"{} // comment"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({}); + assert_eq!(actual, expected); + + // Should not remove comments inside strings + let fixture = r#""/* foo */""#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("/* foo */"); + assert_eq!(actual, expected); +} + +#[test] +fn test_unicode_support() { + // Unicode characters in strings + let fixture = r#""★""#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("★"); + assert_eq!(actual, expected); + + let fixture = r#""😀""#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("😀"); + assert_eq!(actual, expected); + + // Escaped unicode + let fixture = r#""\\u2605""#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("\\u2605"); + assert_eq!(actual, expected); + + // Unicode in keys + let fixture = r#"{"★":true}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"★": true}); + assert_eq!(actual, expected); +} diff --git a/crates/forge_json_repair/tests/complex_repairs.rs b/crates/forge_json_repair/tests/complex_repairs.rs new file mode 100644 index 0000000000000000000000000000000000000000..9e84df78ea12ec23c13857685d221a5f28e26cce --- /dev/null +++ b/crates/forge_json_repair/tests/complex_repairs.rs @@ -0,0 +1,224 @@ +use forge_json_repair::json_repair; +use pretty_assertions::assert_eq; + +#[test] +fn test_newline_separated_json() { + // Basic newline separated JSON (like MongoDB output) + let fixture = "/* 1 */\n{}\n\n/* 2 */\n{}\n\n/* 3 */\n{}\n"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!([{}, {}, {}]); + assert_eq!(actual, expected); + + // With existing commas + let fixture = "/* 1 */\n{},\n\n/* 2 */\n{},\n\n/* 3 */\n{}\n"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!([{}, {}, {}]); + assert_eq!(actual, expected); + + // With trailing comma + let fixture = "/* 1 */\n{},\n\n/* 2 */\n{},\n\n/* 3 */\n{},\n"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!([{}, {}, {}]); + assert_eq!(actual, expected); +} + +#[test] +fn test_comma_separated_lists() { + let fixture = "1,2,3"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!([1, 2, 3]); + assert_eq!(actual, expected); + + let fixture = "1,2,3,"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!([1, 2, 3]); + assert_eq!(actual, expected); + + let fixture = "1\n2\n3"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!([1, 2, 3]); + assert_eq!(actual, expected); + + let fixture = "a\nb"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!(["a", "b"]); + assert_eq!(actual, expected); + + let fixture = "a,b"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!(["a", "b"]); + assert_eq!(actual, expected); +} + +#[test] +fn test_missing_comma_repairs() { + // Missing commas between array items + let fixture = r#"{"array": [{}{}]}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"array": [{}, {}]}); + assert_eq!(actual, expected); + + let fixture = r#"{"array": [{} {}]}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"array": [{}, {}]}); + assert_eq!(actual, expected); + + let fixture = "{\n\"array\": [{}\n{}]\n}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"array": [{}, {}]}); + assert_eq!(actual, expected); + + let fixture = "{\n\"array\": [\n{}\n{}\n]\n}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"array": [{}, {}]}); + assert_eq!(actual, expected); + + let fixture = "{\n\"array\": [\n1\n2\n]\n}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"array": [1, 2]}); + assert_eq!(actual, expected); + + let fixture = "{\n\"array\": [\n\"a\"\n\"b\"\n]\n}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"array": ["a", "b"]}); + assert_eq!(actual, expected); + + // Should leave normal array as is + let fixture = "[\n{},\n{}\n]"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!([{}, {}]); + assert_eq!(actual, expected); +} + +#[test] +fn test_missing_comma_between_object_properties() { + let fixture = "{\"a\":2\n\"b\":3\n}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": 2, "b": 3}); + assert_eq!(actual, expected); + + let fixture = "{\"a\":2\n\"b\":3\n\"c\":4}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": 2, "b": 3, "c": 4}); + assert_eq!(actual, expected); + + let fixture = "{\n \"firstName\": \"John\"\n \"lastName\": \"Smith\"}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"firstName": "John", "lastName": "Smith"}); + assert_eq!(actual, expected); + + let fixture = "{\n \"firstName\": \"John\" /* comment */ \n \"lastName\": \"Smith\"}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"firstName": "John", "lastName": "Smith"}); + assert_eq!(actual, expected); + + // Verify parsing a comma after a return + let fixture = "{\n \"firstName\": \"John\"\n , \"lastName\": \"Smith\"}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"firstName": "John", "lastName": "Smith"}); + assert_eq!(actual, expected); +} + +#[test] +fn test_repair_combination_of_issues() { + let fixture = "{\n\"array\": [\na\nb\n]\n}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"array": ["a", "b"]}); + assert_eq!(actual, expected); + + let fixture = "1\n2"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!([1, 2]); + assert_eq!(actual, expected); + + let fixture = "[a,b\nc]"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!(["a", "b", "c"]); + assert_eq!(actual, expected); +} + +#[test] +fn test_repair_numbers_at_end() { + let fixture = r#"{"a":2.}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": 2.0}); + assert_eq!(actual, expected); + + let fixture = r#"{"a":2e}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": 2e0}); + assert_eq!(actual, expected); + + let fixture = r#"{"a":2e-}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": 2e0}); + assert_eq!(actual, expected); + + let fixture = r#"{"a":-}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": -0.0}); + assert_eq!(actual, expected); + + let fixture = "[2e,]"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!([2e0]); + assert_eq!(actual, expected); + + let fixture = "[2e ]"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!([2e0]); + assert_eq!(actual, expected); + + let fixture = "[-,]"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!([-0.0]); + assert_eq!(actual, expected); +} + +#[test] +fn test_remove_redundant_closing_brackets() { + let fixture = r#"{"a": 1}}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": 1}); + assert_eq!(actual, expected); + + let fixture = r#"{"a": 1}}]}}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": 1}); + assert_eq!(actual, expected); + + let fixture = r#"{"a": 1 } } ] } "#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": 1}); + assert_eq!(actual, expected); + + let fixture = r#"{"a":2]"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": 2}); + assert_eq!(actual, expected); + + let fixture = r#"{"a":2,]"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": 2}); + assert_eq!(actual, expected); + + let fixture = "{}}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({}); + assert_eq!(actual, expected); + + let fixture = "[2,}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!([2]); + assert_eq!(actual, expected); + + let fixture = "[}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!([]); + assert_eq!(actual, expected); + + let fixture = "{]"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({}); + assert_eq!(actual, expected); +} diff --git a/crates/forge_json_repair/tests/error_cases.rs b/crates/forge_json_repair/tests/error_cases.rs new file mode 100644 index 0000000000000000000000000000000000000000..5f8d6fefc55dba1b8959fcf4d0ede8db617f71fa --- /dev/null +++ b/crates/forge_json_repair/tests/error_cases.rs @@ -0,0 +1,96 @@ +use forge_json_repair::json_repair; +use pretty_assertions::assert_eq; + +#[test] +fn test_error_cases() { + // Empty string + assert!(json_repair::("").is_err()); + + // Missing colon + assert!(json_repair::(r#"{"a","#).is_err()); + + // Missing object key + assert!(json_repair::("{:2}").is_err()); + + // Unexpected character after valid JSON + assert!(json_repair::(r#"{"a":2}{}"#).is_err()); + + // Invalid unicode + assert!(json_repair::(r#""\u26""#).is_err()); + assert!(json_repair::(r#""\uZ000""#).is_err()); +} + +#[test] +fn test_regex_single_slash() { + // This test case triggers index out of bounds at line 765 and 771 in + // parse_regex When self.i == 0, accessing self.chars.get(self.i - 1) causes + // underflow After processing single '/', self.i becomes 2 but chars only + // has length 1 + let fixture = "/"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("/"); + assert_eq!(actual, expected); +} + +#[test] +fn test_regex_with_backslash_slash() { + // Test regex with escaped slash at the end + // parse_regex treats the regex as a string literal, so backslash is preserved + let fixture = r#"/a\/"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!(r#"/a\/"#); + assert_eq!(actual, expected); +} + +#[test] +fn test_string_with_colon_at_start() { + // This test case checks for potential index out of bounds at line 445 + // When self.i == 0 and we try to access self.chars.get(self.i - 1) + let fixture = ":"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!(":"); + assert_eq!(actual, expected); +} + +#[test] +fn test_multibyte_unicode_missing_end_quote() { + // Triggers index out of bounds in insert_before_last_whitespace_str. + // The output buffer contains multi-byte UTF-8 characters (é = 2 bytes each), + // so self.output.len() (byte count) > chars.len() (char count). + // When the repair path calls insert_before_last_whitespace_str with trailing + // whitespace, it initialises `index` from the byte length and then indexes + // into a Vec at that byte-length position, panicking. + let fixture = r#""café "#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("café"); + assert_eq!(actual, expected); +} + +#[test] +fn test_multibyte_unicode_missing_comma_in_object() { + // Triggers index out of bounds in insert_before_last_whitespace_str (line 459). + // parse_string first collects `"é,"` and hits the inner `"test"`. The + // prev_non_whitespace char is `,`, so it retries with stop_at_index=2 + // (the comma position). On retry it collects str_content = `"é` (3 bytes, + // 2 chars) and hits stop_at_index, calling insert_before_last_whitespace_str. + // That function sets index = text.len() = 3 (byte count) and then accesses + // chars[index - 1] = chars[2] on a Vec of length 2 — panic. + let fixture = "\"é,\"test\""; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!(["é", "test"]); + assert_eq!(actual, expected); +} + +#[test] +fn test_multibyte_unicode_missing_closing_brace() { + // Triggers index out of bounds in insert_before_last_whitespace_str (line 384). + // A string with a multi-byte character followed by trailing whitespace and + // no closing quote hits the "end of text, missing end quote" repair path. + // str_content = `"🎉 ` (6 bytes, 3 chars). insert_before_last_whitespace_str + // sets index = text.len() = 6 and accesses chars[5] on a Vec of + // length 3 — panic. + let fixture = "\"🎉 "; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("🎉"); + assert_eq!(actual, expected); +} diff --git a/crates/forge_json_repair/tests/escaping.rs b/crates/forge_json_repair/tests/escaping.rs new file mode 100644 index 0000000000000000000000000000000000000000..970bd565da5d820089631446715a71aeb195f266 --- /dev/null +++ b/crates/forge_json_repair/tests/escaping.rs @@ -0,0 +1,56 @@ +use forge_json_repair::json_repair; +use pretty_assertions::assert_eq; + +#[test] +fn test_escape_characters() { + let fixture = r#""foo'bar""#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("foo'bar"); + assert_eq!(actual, expected); + + let fixture = r#""foo\"bar""#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("foo\"bar"); + assert_eq!(actual, expected); + + let fixture = "'foo\"bar'"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("foo\"bar"); + assert_eq!(actual, expected); + + let fixture = "'foo\\'bar'"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("foo'bar"); + assert_eq!(actual, expected); +} + +#[test] +fn test_escape_control_characters() { + let fixture = "\"hello\nworld\""; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("hello\nworld"); + assert_eq!(actual, expected); + + let fixture = "\"hello\tworld\""; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("hello\tworld"); + assert_eq!(actual, expected); + + let fixture = "\"hello\rworld\""; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("hello\rworld"); + assert_eq!(actual, expected); +} + +#[test] +fn test_escape_unescaped_quotes() { + let fixture = r#""The TV has a 24" screen""#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("The TV has a 24\" screen"); + assert_eq!(actual, expected); + + let fixture = r#"{"key": "apple "bee" carrot"}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"key": "apple \"bee\" carrot"}); + assert_eq!(actual, expected); +} diff --git a/crates/forge_json_repair/tests/numbers.rs b/crates/forge_json_repair/tests/numbers.rs new file mode 100644 index 0000000000000000000000000000000000000000..609cf739d8cdfafe7665f97755a5518fd0fdf42f --- /dev/null +++ b/crates/forge_json_repair/tests/numbers.rs @@ -0,0 +1,58 @@ +use forge_json_repair::json_repair; +use pretty_assertions::assert_eq; + +#[test] +fn test_invalid_numbers_to_strings() { + let fixture = "ES2020"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("ES2020"); + assert_eq!(actual, expected); + + let fixture = "0.0.1"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("0.0.1"); + assert_eq!(actual, expected); + + let fixture = "746de9ad-d4ff-4c66-97d7-00a92ad46967"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("746de9ad-d4ff-4c66-97d7-00a92ad46967"); + assert_eq!(actual, expected); + + let fixture = "234..5"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("234..5"); + assert_eq!(actual, expected); + + let fixture = "2e3.4"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("2e3.4"); + assert_eq!(actual, expected); +} + +#[test] +fn test_numbers_with_leading_zeros() { + let fixture = "0789"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("0789"); + assert_eq!(actual, expected); + + let fixture = "000789"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("000789"); + assert_eq!(actual, expected); + + let fixture = "001.2"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("001.2"); + assert_eq!(actual, expected); + + let fixture = "[0789]"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!(["0789"]); + assert_eq!(actual, expected); + + let fixture = "{value:0789}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"value": "0789"}); + assert_eq!(actual, expected); +} diff --git a/crates/forge_json_repair/tests/punctuation.rs b/crates/forge_json_repair/tests/punctuation.rs new file mode 100644 index 0000000000000000000000000000000000000000..4a5dfcb63bcc5e5676de3b83531f4db9a4424b4d --- /dev/null +++ b/crates/forge_json_repair/tests/punctuation.rs @@ -0,0 +1,81 @@ +use forge_json_repair::json_repair; +use pretty_assertions::assert_eq; + +#[test] +fn test_repair_missing_commas() { + // Array items + let fixture = r#"{"array": [{}{}]}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"array": [{}, {}]}); + assert_eq!(actual, expected); + + // Object properties + let fixture = r#"{"a":2"b":3}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": 2, "b": 3}); + assert_eq!(actual, expected); +} + +#[test] +fn test_repair_missing_colons() { + let fixture = r#"{"a" "b"}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": "b"}); + assert_eq!(actual, expected); + + let fixture = r#"{"a" 2}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": 2}); + assert_eq!(actual, expected); + + let fixture = r#"{"a" true}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": true}); + assert_eq!(actual, expected); + + let fixture = r#"{"a" false}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": false}); + assert_eq!(actual, expected); + + let fixture = r#"{"a" null}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": null}); + assert_eq!(actual, expected); +} + +#[test] +fn test_strip_leading_commas() { + let fixture = "[,1,2,3]"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!([1, 2, 3]); + assert_eq!(actual, expected); + + let fixture = r#"{,"message": "hi"}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"message": "hi"}); + assert_eq!(actual, expected); +} + +#[test] +fn test_strip_trailing_commas() { + let fixture = "[1,2,3,]"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!([1, 2, 3]); + assert_eq!(actual, expected); + + let fixture = r#"{"a":2,}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": 2}); + assert_eq!(actual, expected); + + let fixture = "4,"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!(4); + assert_eq!(actual, expected); + + let fixture = r#"{"a":2},"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": 2}); + assert_eq!(actual, expected); +} diff --git a/crates/forge_json_repair/tests/quotes.rs b/crates/forge_json_repair/tests/quotes.rs new file mode 100644 index 0000000000000000000000000000000000000000..feafdce5c68c6610b6ec13724a43499fe96e0cd9 --- /dev/null +++ b/crates/forge_json_repair/tests/quotes.rs @@ -0,0 +1,112 @@ +use forge_json_repair::json_repair; +use pretty_assertions::assert_eq; + +#[test] +fn test_add_missing_quotes() { + // Simple unquoted strings + let fixture = "abc"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("abc"); + assert_eq!(actual, expected); + + let fixture = "hello world"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("hello world"); + assert_eq!(actual, expected); + + // Unquoted object keys + let fixture = "{a:2}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": 2}); + assert_eq!(actual, expected); + + let fixture = "{a: 2}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": 2}); + assert_eq!(actual, expected); + + let fixture = "{2: 2}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"2": 2}); + assert_eq!(actual, expected); + + let fixture = "{true: 2}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"true": 2}); + assert_eq!(actual, expected); + + // Unquoted array values + let fixture = "[a,b]"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!(["a", "b"]); + assert_eq!(actual, expected); +} + +#[test] +fn test_add_missing_end_quote() { + let fixture = r#""abc"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("abc"); + assert_eq!(actual, expected); + + let fixture = "'abc"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("abc"); + assert_eq!(actual, expected); + + let fixture = r#"{"a":"b}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": "b"}); + assert_eq!(actual, expected); + + let fixture = r#"{"a":"b,"c":"d"}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": "b", "c": "d"}); + assert_eq!(actual, expected); +} + +#[test] +fn test_add_missing_start_quote() { + let fixture = r#"abc""#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("abc"); + assert_eq!(actual, expected); + + let fixture = r#"[a","b"]"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!(["a", "b"]); + assert_eq!(actual, expected); + + let fixture = r#"{"a":"foo","b":"bar"}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": "foo", "b": "bar"}); + assert_eq!(actual, expected); + + let fixture = r#"{a":"foo","b":"bar"}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": "foo", "b": "bar"}); + assert_eq!(actual, expected); +} + +#[test] +fn test_replace_single_quotes() { + let fixture = "{'a':2}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": 2}); + assert_eq!(actual, expected); + + let fixture = "{'a':'foo'}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": "foo"}); + assert_eq!(actual, expected); + + let fixture = r#"{"a":'foo'}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": "foo"}); + assert_eq!(actual, expected); + + let fixture = "{a:'foo',b:'bar'}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": "foo", "b": "bar"}); + assert_eq!(actual, expected); +} diff --git a/crates/forge_json_repair/tests/special_quotes.rs b/crates/forge_json_repair/tests/special_quotes.rs new file mode 100644 index 0000000000000000000000000000000000000000..ec6de5ba523b77f4e9d2b5d98734f00cdaf6c589 --- /dev/null +++ b/crates/forge_json_repair/tests/special_quotes.rs @@ -0,0 +1,170 @@ +use forge_json_repair::json_repair; +use pretty_assertions::assert_eq; + +#[test] +fn test_escaped_string_contents() { + let fixture = r#"\"hello world\""#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("hello world"); + assert_eq!(actual, expected); + + let fixture = r#"\"hello world\"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("hello world"); + assert_eq!(actual, expected); + + let fixture = r#"\"hello \"world\"\""#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("hello \"world\""); + assert_eq!(actual, expected); + + let fixture = r#"[\"hello \"world\"\"]"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!(["hello \"world\""]); + assert_eq!(actual, expected); + + let fixture = r#"{\"stringified\": \"hello \"world\"\"}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"stringified": "hello \"world\""}); + assert_eq!(actual, expected); + + // Weird but close to likely intention + let fixture = r#"[\"hello\, \"world\"]"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!(["hello, \"world"]); + assert_eq!(actual, expected); + + // Invalid end quote handling + let fixture = r#"\"hello""#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("hello"); + assert_eq!(actual, expected); +} + +#[test] +fn test_special_quote_characters() { + // Left/right single quotes (using unicode escapes) + let fixture = "\u{2018}foo\u{2019}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("foo"); + assert_eq!(actual, expected); + + // Left/right double quotes (using unicode escapes) + let fixture = "\u{201C}foo\u{201D}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("foo"); + assert_eq!(actual, expected); + + // Backtick and acute accent + let fixture = "\u{0060}foo\u{00B4}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("foo"); + assert_eq!(actual, expected); + + // Mixed quotes + let fixture = "\u{0060}foo'"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("foo"); + assert_eq!(actual, expected); + + // In object syntax + let fixture = "{\u{2018}a\u{2019}:\u{2018}b\u{2019}}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": "b"}); + assert_eq!(actual, expected); + + let fixture = "{\u{0060}a\u{00B4}:\u{0060}b\u{00B4}}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": "b"}); + assert_eq!(actual, expected); +} + +#[test] +fn test_special_quotes_inside_strings() { + // Should not replace special quotes inside normal strings + let fixture = "\"Rounded \u{201C} quote\""; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("Rounded \u{201C} quote"); + assert_eq!(actual, expected); + + let fixture = "'\u{201C}Rounded quote'"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("\u{201C}Rounded quote"); + assert_eq!(actual, expected); + + let fixture = "\"Rounded \u{2018} quote\""; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("Rounded \u{2018} quote"); + assert_eq!(actual, expected); + + let fixture = "'\u{2018}Rounded quote'"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("\u{2018}Rounded quote"); + assert_eq!(actual, expected); + + let fixture = r#"'Double " quote'"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("Double \" quote"); + assert_eq!(actual, expected); +} + +#[test] +fn test_quote_repair_edge_cases() { + // Should leave string content untouched + let fixture = r#""{a:b}""#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("{a:b}"); + assert_eq!(actual, expected); +} + +#[test] +fn test_special_whitespace_characters() { + // Non-breaking space and other special spaces + let fixture = "{\"a\":\u{00a0}\"foo\u{00a0}bar\"}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": "foo\u{00a0}bar"}); + assert_eq!(actual, expected); + + let fixture = "{\"a\":\u{202F}\"foo\"}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": "foo"}); + assert_eq!(actual, expected); + + let fixture = "{\"a\":\u{205F}\"foo\"}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": "foo"}); + assert_eq!(actual, expected); + + let fixture = "{\"a\":\u{3000}\"foo\"}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": "foo"}); + assert_eq!(actual, expected); +} + +#[test] +fn test_stop_at_newline_for_missing_quotes() { + let fixture = "[\n\"abc,\n\"def\"\n]"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!(["abc", "def"]); + assert_eq!(actual, expected); + + let fixture = "[\n\"abc, \n\"def\"\n]"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!(["abc", "def"]); + assert_eq!(actual, expected); + + let fixture = "[\"abc]\n"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!(["abc"]); + assert_eq!(actual, expected); + + let fixture = "[\"abc ]\n"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!(["abc"]); + assert_eq!(actual, expected); + + let fixture = "[\n[\n\"abc\n]\n]\n"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!([["abc"]]); + assert_eq!(actual, expected); +} diff --git a/crates/forge_json_repair/tests/special_values.rs b/crates/forge_json_repair/tests/special_values.rs new file mode 100644 index 0000000000000000000000000000000000000000..98b7b34b0be25a61ff041e444193e359b66af72b --- /dev/null +++ b/crates/forge_json_repair/tests/special_values.rs @@ -0,0 +1,56 @@ +use forge_json_repair::json_repair; +use pretty_assertions::assert_eq; + +#[test] +fn test_python_constants() { + let fixture = "True"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!(true); + assert_eq!(actual, expected); + + let fixture = "False"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!(false); + assert_eq!(actual, expected); + + let fixture = "None"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!(null); + assert_eq!(actual, expected); +} + +#[test] +fn test_repair_undefined_values() { + let fixture = r#"{"a":undefined}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": null}); + assert_eq!(actual, expected); + + let fixture = "[undefined]"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!([null]); + assert_eq!(actual, expected); + + let fixture = "undefined"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!(null); + assert_eq!(actual, expected); +} + +#[test] +fn test_repair_missing_object_value() { + let fixture = r#"{"a":}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": null}); + assert_eq!(actual, expected); + + let fixture = r#"{"a":,"b":2}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": null, "b": 2}); + assert_eq!(actual, expected); + + let fixture = r#"{"a":"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": null}); + assert_eq!(actual, expected); +} diff --git a/crates/forge_json_repair/tests/strings.rs b/crates/forge_json_repair/tests/strings.rs new file mode 100644 index 0000000000000000000000000000000000000000..c59b2f054dd176166070984d8f15ee99b7f92847 --- /dev/null +++ b/crates/forge_json_repair/tests/strings.rs @@ -0,0 +1,61 @@ +use forge_json_repair::json_repair; +use pretty_assertions::assert_eq; + +#[test] +fn test_parse_unquoted_strings() { + let fixture = "hello world"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("hello world"); + assert_eq!(actual, expected); + + let fixture = "She said: no way"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("She said: no way"); + assert_eq!(actual, expected); +} + +#[test] +fn test_turn_symbols_into_strings() { + let fixture = "foo"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("foo"); + assert_eq!(actual, expected); + + let fixture = "[1,foo,4]"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!([1, "foo", 4]); + assert_eq!(actual, expected); + + let fixture = "{foo: bar}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"foo": "bar"}); + assert_eq!(actual, expected); + + let fixture = "foo 2 bar"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("foo 2 bar"); + assert_eq!(actual, expected); + + let fixture = "{greeting: hello world}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"greeting": "hello world"}); + assert_eq!(actual, expected); +} + +#[test] +fn test_repair_urls() { + let fixture = "https://www.example.com/"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("https://www.example.com/"); + assert_eq!(actual, expected); + + let fixture = "{url:https://www.example.com/}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"url": "https://www.example.com/"}); + assert_eq!(actual, expected); + + let fixture = r#"{"url":"https://www.example.com/}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"url": "https://www.example.com/"}); + assert_eq!(actual, expected); +} diff --git a/crates/forge_json_repair/tests/truncation.rs b/crates/forge_json_repair/tests/truncation.rs new file mode 100644 index 0000000000000000000000000000000000000000..0b79e11f7c46d9504fc7de9e46a78b13c9e56b99 --- /dev/null +++ b/crates/forge_json_repair/tests/truncation.rs @@ -0,0 +1,148 @@ +use forge_json_repair::json_repair; +use pretty_assertions::assert_eq; + +#[test] +fn test_add_missing_closing_brackets() { + let fixture = "{"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({}); + assert_eq!(actual, expected); + + let fixture = r#"{"a":2"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": 2}); + assert_eq!(actual, expected); + + let fixture = r#"{"a":2,"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": 2}); + assert_eq!(actual, expected); + + let fixture = "["; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!([]); + assert_eq!(actual, expected); + + let fixture = "[1,2,3"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!([1, 2, 3]); + assert_eq!(actual, expected); + + let fixture = "[1,2,3,"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!([1, 2, 3]); + assert_eq!(actual, expected); +} + +#[test] +fn test_repair_truncated_json() { + let fixture = r#""foo"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!("foo"); + assert_eq!(actual, expected); + + let fixture = "["; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!([]); + assert_eq!(actual, expected); + + let fixture = r#"["foo"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!(["foo"]); + assert_eq!(actual, expected); + + let fixture = r#"["foo","#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!(["foo"]); + assert_eq!(actual, expected); + + let fixture = r#"{"foo":"bar""#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"foo": "bar"}); + assert_eq!(actual, expected); + + let fixture = r#"{"foo":"bar"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"foo": "bar"}); + assert_eq!(actual, expected); + + let fixture = r#"{"foo":"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"foo": null}); + assert_eq!(actual, expected); + + let fixture = r#"{"foo""#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"foo": null}); + assert_eq!(actual, expected); + + let fixture = r#"{"foo"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"foo": null}); + assert_eq!(actual, expected); + + let fixture = "{"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({}); + assert_eq!(actual, expected); + + let fixture = "2."; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!(2.0); + assert_eq!(actual, expected); + + let fixture = "2e"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!(2e0); + assert_eq!(actual, expected); + + let fixture = "2e+"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!(2e0); + assert_eq!(actual, expected); + + let fixture = "2e-"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!(2e0); + assert_eq!(actual, expected); +} + +#[test] +fn test_repair_ellipsis() { + // Arrays + let fixture = "[1,2,3,...]"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!([1, 2, 3]); + assert_eq!(actual, expected); + + let fixture = "[1,2,3,...,9]"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!([1, 2, 3, 9]); + assert_eq!(actual, expected); + + let fixture = "[...,7,8,9]"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!([7, 8, 9]); + assert_eq!(actual, expected); + + let fixture = "[...]"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!([]); + assert_eq!(actual, expected); + + // Objects + let fixture = r#"{"a":2,"b":3,...}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": 2, "b": 3}); + assert_eq!(actual, expected); + + let fixture = r#"{"a":2,"b":3,...,"z":26}"#; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({"a": 2, "b": 3, "z": 26}); + assert_eq!(actual, expected); + + let fixture = "{...}"; + let actual = json_repair::(fixture).unwrap(); + let expected = serde_json::json!({}); + assert_eq!(actual, expected); +} diff --git a/crates/forge_json_repair/tests/valid_json.rs b/crates/forge_json_repair/tests/valid_json.rs new file mode 100644 index 0000000000000000000000000000000000000000..1c731170153325e9a98747e5044a0a1c295e7001 --- /dev/null +++ b/crates/forge_json_repair/tests/valid_json.rs @@ -0,0 +1,71 @@ +use forge_json_repair::json_repair; +use pretty_assertions::assert_eq; + +#[test] +fn test_parse_valid_json() { + // Full JSON object + let text = r#"{"a":2.3e100,"b":"str","c":null,"d":false,"e":[1,2,3]}"#; + let actual = json_repair::(text).unwrap(); + let expected = serde_json::from_str::(text).unwrap(); + assert_eq!(actual, expected); + + // Whitespace + assert_repair(" { \n } \t "); + + // Objects + assert_repair("{}"); + assert_repair("{ }"); + assert_repair(r#"{"a": {}}"#); + assert_repair(r#"{"a": "b"}"#); + assert_repair(r#"{"a": 2}"#); + + // Arrays + assert_repair("[]"); + assert_repair("[ ]"); + assert_repair("[1,2,3]"); + assert_repair("[ 1 , 2 , 3 ]"); + assert_repair("[1,2,[3,4,5]]"); + assert_repair("[{}]"); + assert_repair(r#"{"a":[]}"#); + assert_repair(r#"[1, "hi", true, false, null, {}, []]"#); + + // Numbers + assert_repair("23"); + assert_repair("0"); + assert_repair("0e+2"); + assert_repair("0.0"); + assert_repair("-0"); + assert_repair("2.3"); + assert_repair("2300e3"); + assert_repair("2300e+3"); + assert_repair("2300e-3"); + assert_repair("-2"); + assert_repair("2e-3"); + assert_repair("2.3e-3"); + + // Strings + assert_repair(r#""str""#); + assert_repair(r#""\"\\\\/\\b\\f\\n\\r\\t""#); + assert_repair(r#""\\u260E""#); + + // Keywords + assert_repair("true"); + assert_repair("false"); + assert_repair("null"); + + // String delimiters + assert_repair(r#""""#); + assert_repair(r#""[""#); + assert_repair(r#""]""#); + assert_repair(r#""{""#); + assert_repair(r#""}""#); + assert_repair(r#"":""#); + assert_repair(r#"",""#); +} + +// Helper function to assert that input is repaired to itself (valid JSON) +fn assert_repair(text: &str) { + let actual = json_repair::(text).unwrap(); + let expected = serde_json::from_str::(text).unwrap(); + assert_eq!(actual, expected); +} diff --git a/crates/forge_main/src/porcelain.rs b/crates/forge_main/src/porcelain.rs new file mode 100644 index 0000000000000000000000000000000000000000..6454d1c206d150c5873ff6d3ae8d8a1e2c90ef2a --- /dev/null +++ b/crates/forge_main/src/porcelain.rs @@ -0,0 +1,1248 @@ +use std::collections::HashMap; +use std::fmt; + +use indexmap::IndexSet; + +use crate::display_constants::headers; +use crate::info::{Info, Section}; + +/// Porcelain is an intermediate representation that converts Info into a flat, +/// tabular structure suitable for machine-readable output. +/// +/// Structure: Vec<(String, Vec>)> +/// - First element: Section name +/// - Second element: Vec of Option pairs where: +/// - Index 0, 2, 4... are keys +/// - Index 1, 3, 5... are values +/// - None = missing value +#[derive(Debug, PartialEq, Clone)] +pub struct Porcelain(Vec>>); + +impl Porcelain { + /// Creates a new empty Porcelain instance + pub fn new() -> Self { + Porcelain(Vec::new()) + } + + /// Skips the first n rows + pub fn skip(self, n: usize) -> Self { + Porcelain(self.0.into_iter().skip(n).collect()) + } + + #[allow(unused)] + pub fn drop_col(self, c: usize) -> Self { + Porcelain( + self.0 + .into_iter() + .map(|row| { + row.into_iter() + .enumerate() + .filter_map(|(i, col)| if i == c { None } else { Some(col) }) + .collect() + }) + .collect(), + ) + } + + /// Drops multiple columns at once + /// + /// # Arguments + /// * `cols` - A slice of column indices to drop + /// + /// # Example + /// ```rust,ignore + /// let porcelain = Porcelain::new(); + /// let result = porcelain.drop_cols(&[0, 2, 4]); + /// ``` + #[allow(unused)] + pub fn drop_cols(self, cols: &[usize]) -> Self { + Porcelain( + self.0 + .into_iter() + .map(|row| { + row.into_iter() + .enumerate() + .filter_map(|(i, col)| if cols.contains(&i) { None } else { Some(col) }) + .collect() + }) + .collect(), + ) + } + + /// Maps a function over all cells in the specified column + pub fn map_col(self, c: usize, f: F) -> Self + where + F: Fn(Option) -> Option, + { + Porcelain( + self.0 + .into_iter() + .map(|row| { + row.into_iter() + .enumerate() + .map(|(i, col)| if i == c { f(col) } else { col }) + .collect() + }) + .collect(), + ) + } + + /// Truncates the specified column to a maximum number of characters, + /// appending "..." after the kept characters if truncated. The "..." is + /// not counted toward `max_len`. + pub fn truncate(self, c: usize, max_len: usize) -> Self { + self.map_col(c, |col| { + col.map(|value| { + if value.chars().count() > max_len { + let truncated: String = value.chars().take(max_len).collect(); + format!("{}...", truncated) + } else { + value + } + }) + }) + } + + pub fn swap_cols(self, col1: usize, col2: usize) -> Self { + Porcelain( + self.0 + .into_iter() + .map(|mut row| { + if row.len() > col1.max(col2) { + row.swap(col1, col2); + } + row + }) + .collect(), + ) + } + + /// Sorts rows based on multiple columns + /// + /// Preserves the header row (first row) and sorts all subsequent rows. + /// Columns are sorted in the order specified in `cols`. None values are + /// sorted after Some values. + /// + /// # Arguments + /// * `cols` - Column indices to sort by, in order of precedence + /// + /// # Example + /// ```ignore + /// porcelain.sort_by(&[1, 2]) // Sort by column 1, then by column 2 + /// ``` + pub fn sort_by(self, cols: &[usize]) -> Self { + if self.0.is_empty() || cols.is_empty() { + return self; + } + + let mut rows = self.0; + let header = if !rows.is_empty() { + Some(rows.remove(0)) + } else { + None + }; + + rows.sort_by(|a, b| { + for &col in cols { + let a_val = a.get(col); + let b_val = b.get(col); + + let ordering = match (a_val, b_val) { + (Some(Some(a)), Some(Some(b))) => a.cmp(b), + (Some(Some(_)), Some(None)) => std::cmp::Ordering::Less, + (Some(None), Some(Some(_))) => std::cmp::Ordering::Greater, + (Some(None), Some(None)) => std::cmp::Ordering::Equal, + (Some(_), None) => std::cmp::Ordering::Less, + (None, Some(_)) => std::cmp::Ordering::Greater, + (None, None) => std::cmp::Ordering::Equal, + }; + + if ordering != std::cmp::Ordering::Equal { + return ordering; + } + } + std::cmp::Ordering::Equal + }); + + if let Some(header) = header { + rows.insert(0, header); + } + + Porcelain(rows) + } + + /// Applies case transformation to specified columns + /// + /// # Arguments + /// * `cols` - Column indices to transform + /// * `case` - The case to apply (e.g., Case::Snake, Case::Kebab) + /// + /// # Example + /// ```ignore + /// use convert_case::Case; + /// + /// porcelain.to_case(&[0, 1], Case::Snake) + /// ``` + #[allow(clippy::wrong_self_convention)] + pub fn to_case(self, cols: &[usize], case: convert_case::Case) -> Self { + use convert_case::Casing; + + Porcelain( + self.0 + .into_iter() + .map(|row| { + row.into_iter() + .enumerate() + .map(|(i, col)| { + if cols.contains(&i) { + col.map(|v| v.to_case(case)) + } else { + col + } + }) + .collect() + }) + .collect(), + ) + } + + #[allow(unused)] + pub fn into_body(self) -> Vec>> { + // Skip headers and return + self.0.into_iter().skip(1).collect() + } + + #[allow(unused)] + pub fn into_rows(self) -> Vec>> { + self.0 + } + + /// Converts from wide format to long format. + /// + /// Transforms entity-centric rows (wide format with many columns) into + /// field-centric rows (long format with three columns: entity_id, + /// field_name, field_value). This is also known as unpivoting or + /// melting in data transformation terminology. + /// + /// # Example + /// Input (wide format): + /// ```text + /// Headers: [$ID, version, shell, id, title, model] + /// Row 1: [env, 0.1.0, zsh, None, None, None] + /// Row 2: [conversation, None, None, 000-000-000, make agents great again, None] + /// ``` + /// + /// Output (long format): + /// ```text + /// Headers: [$ID, field, value] + /// Row 1: [env, version, 0.1.0] + /// Row 2: [env, shell, zsh] + /// Row 3: [conversation, id, 000-000-000] + /// Row 4: [conversation, title, make agents great again] + /// ``` + /// Converts all headers (first row) to uppercase + pub fn uppercase_headers(self) -> Self { + if self.0.is_empty() { + return self; + } + + let mut rows = self.0; + if let Some(header_row) = rows.first_mut() { + *header_row = header_row + .iter() + .map(|col| col.as_ref().map(|s| s.to_uppercase())) + .collect(); + } + + Porcelain(rows) + } + + /// Sets the headers (first row) of the Porcelain structure + /// + /// Replaces the first row with the provided headers. If the Porcelain is + /// empty, creates a new header row. If an empty iterator is provided, + /// returns self unchanged. + /// + /// # Arguments + /// * `headers` - An iterator of header names to set + pub fn set_headers(self, headers: I) -> Self + where + I: IntoIterator, + S: Into, + { + let new_headers: Vec> = + headers.into_iter().map(|h| Some(h.into())).collect(); + + if new_headers.is_empty() { + return self; + } + + let mut rows = self.0; + + if rows.is_empty() { + // If Porcelain is empty, just add the headers + rows.push(new_headers); + } else { + // Replace the first row with new headers + if let Some(first) = rows.first_mut() { + *first = new_headers; + } + } + + Porcelain(rows) + } + + pub fn into_long(self) -> Self { + if self.0.is_empty() { + return self; + } + + let Some(headers) = self.0.first() else { + return self; + }; + let data_rows = self.0.get(1..).unwrap_or(&[]); + + if data_rows.is_empty() || headers.is_empty() { + return self; + } + + // Create new headers: [$ID, $FIELD, $VALUE] + let new_headers = vec![ + headers + .first() + .cloned() + .unwrap_or(Some(headers::ID.to_string())), + Some(headers::FIELD.to_string()), + Some(headers::VALUE.to_string()), + ]; + + // Create new rows: one row per non-None field for each entity + let mut new_rows = Vec::new(); + + for data_row in data_rows { + // Get the entity ID (first column value) + let entity_id = data_row.first().and_then(|v| v.clone()); + + // For each field in this entity (excluding $ID column) + for (i, value) in data_row.iter().enumerate().skip(1) { + if let Some(value) = value { + let field_name = headers.get(i).and_then(|h| h.clone()); + new_rows.push(vec![entity_id.clone(), field_name, Some(value.to_owned())]); + } + } + } + + let mut result = vec![new_headers]; + result.extend(new_rows); + + Porcelain(result) + } +} + +impl fmt::Display for Porcelain { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + if self.0.is_empty() { + return Ok(()); + } + + // Calculate the maximum width for each column + let num_cols = self.0.iter().map(|row| row.len()).max().unwrap_or(0); + let mut col_widths = vec![0; num_cols]; + + for row in &self.0 { + for (i, cell) in row.iter().enumerate() { + let width = cell.as_ref().map(|s| s.len()).unwrap_or(0); + if let Some(col_width) = col_widths.get_mut(i) { + *col_width = (*col_width).max(width); + } + } + } + + // Format each row + let mut lines = Vec::new(); + for row in &self.0 { + let mut line = String::new(); + for (i, cell) in row.iter().enumerate() { + let content = cell.as_ref().map(|s| s.as_str()).unwrap_or(""); + + if i == row.len() - 1 { + // Last column: no padding + line.push_str(content); + } else { + // Pad to column width + let width = col_widths.get(i).copied().unwrap_or(0); + line.push_str(&format!("{: Self { + Self::new() + } +} + +/// Converts Info to Porcelain representation +/// Handles both cases: +/// - Info with titles: Each title becomes a row with its associated items +/// - Info without titles: Each item becomes its own row +impl From for Porcelain { + fn from(info: Info) -> Self { + Porcelain::from(&info) + } +} + +/// Converts Info reference to Porcelain representation +impl From<&Info> for Porcelain { + fn from(info: &Info) -> Self { + let mut rows = Vec::new(); + let mut cells = HashMap::>::new(); + let mut in_row = false; + // Extract all unique keys + let mut keys = IndexSet::new(); + // Track count of unnamed values separately + let mut value_counter = 1; + let mut last_key: Option = None; + + for section in info.sections() { + match section { + Section::Title(title) => { + if in_row { + rows.push(cells.clone()); + cells = HashMap::new(); + value_counter = 1; + } + + in_row = true; + cells.insert(headers::ID.to_owned(), vec![title.to_owned()]); + keys.insert(headers::ID.to_owned()); + } + Section::Items(key, value) => { + let key = if let Some(key) = key.as_ref().cloned().or(last_key) { + key + } else { + let default_key = format!("{}_{}", headers::VALUE, value_counter); + value_counter += 1; + default_key + }; + last_key = Some(key.clone()); + + cells + .entry(key.to_string()) + .or_default() + .push(value.clone()); + keys.insert(key.to_string()); + } + } + } + + if in_row { + rows.push(cells.clone()); + } + + // Insert Headers + let mut data = vec![ + keys.iter() + .map(|head| Some((*head).to_owned())) + .collect::>(), + ]; + + // Insert Rows + data.extend(rows.iter().map(|rows| { + keys.iter() + .map(|key| { + rows.get(key).and_then(|value| { + if value.is_empty() { + None + } else { + Some(value.join(", ")) + } + }) + }) + .collect::>>() + })); + Porcelain(data) + } +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::*; + + #[test] + fn test_from_info() { + let info = Info::new() + .add_title("user1") + .add_key_value("name", "Alice") + .add_key_value("age", "30") + .add_title("user2") + .add_key_value("name", "Bob") + .add_key_value("age", "25"); + + let actual = Porcelain::from(info).into_body(); + let expected = vec![ + vec![ + Some("user1".into()), + Some("Alice".into()), + Some("30".into()), + ], + vec![Some("user2".into()), Some("Bob".into()), Some("25".into())], + ]; + + assert_eq!(actual, expected) + } + + #[test] + fn test_from_unordered_info() { + let info = Info::new() + .add_title("user1") + .add_key_value("name", "Alice") + .add_key_value("age", "30") + .add_title("user2") + .add_key_value("age", "25") + .add_key_value("name", "Bob"); + + let actual = Porcelain::from(info).into_body(); + let expected = vec![ + vec![ + Some("user1".into()), + Some("Alice".into()), + Some("30".into()), + ], + vec![Some("user2".into()), Some("Bob".into()), Some("25".into())], + ]; + + assert_eq!(actual, expected) + } + + #[test] + fn test_drop_col() { + let info = Porcelain(vec![ + vec![ + Some("user1".into()), + Some("Alice".into()), + Some("30".into()), + ], + vec![Some("user2".into()), Some("Bob".into()), Some("25".into())], + ]); + + let actual = info.drop_col(1).into_rows(); + + let expected = vec![ + vec![Some("user1".into()), Some("30".into())], + vec![Some("user2".into()), Some("25".into())], + ]; + + assert_eq!(actual, expected) + } + + #[test] + fn test_drop_cols() { + let fixture = Porcelain(vec![ + vec![ + Some("user1".into()), + Some("Alice".into()), + Some("30".into()), + Some("Engineer".into()), + ], + vec![ + Some("user2".into()), + Some("Bob".into()), + Some("25".into()), + Some("Designer".into()), + ], + ]); + + let actual = fixture.drop_cols(&[1, 3]).into_rows(); + + let expected = vec![ + vec![Some("user1".into()), Some("30".into())], + vec![Some("user2".into()), Some("25".into())], + ]; + + assert_eq!(actual, expected) + } + + #[test] + fn test_drop_cols_equivalent_to_chained() { + // Test that drop_cols(&[0, 4, 5]) is equivalent to + // drop_col(0).drop_col(3).drop_col(3) + let fixture = Porcelain(vec![ + vec![ + Some("col0".into()), + Some("col1".into()), + Some("col2".into()), + Some("col3".into()), + Some("col4".into()), + Some("col5".into()), + Some("col6".into()), + ], + vec![ + Some("row2_0".into()), + Some("row2_1".into()), + Some("row2_2".into()), + Some("row2_3".into()), + Some("row2_4".into()), + Some("row2_5".into()), + Some("row2_6".into()), + ], + ]); + + let actual = fixture.clone().drop_cols(&[0, 4, 5]).into_rows(); + let expected = fixture.drop_col(0).drop_col(3).drop_col(3).into_rows(); + + assert_eq!(actual, expected) + } + + #[test] + fn test_map_col() { + let info = Porcelain(vec![ + vec![ + Some("user1".into()), + Some("Alice".into()), + Some("30".into()), + ], + vec![Some("user2".into()), Some("Bob".into()), Some("25".into())], + ]); + + let actual = info + .map_col(1, |col| col.map(|v| v.to_uppercase())) + .into_rows(); + + let expected = vec![ + vec![ + Some("user1".into()), + Some("ALICE".into()), + Some("30".into()), + ], + vec![Some("user2".into()), Some("BOB".into()), Some("25".into())], + ]; + + assert_eq!(actual, expected) + } + #[test] + fn test_truncate() { + let info = Porcelain(vec![ + vec![ + Some("user1".into()), + Some("Alice".into()), + Some("very_long_name".into()), + ], + vec![ + Some("user2".into()), + Some("Bob".into()), + Some("short".into()), + ], + ]); + + // truncate(2, 5): "very_long_name" has 14 chars > 5, keep 5 then append "..." + let actual = info.truncate(2, 5).into_rows(); + + let expected = vec![ + vec![ + Some("user1".into()), + Some("Alice".into()), + Some("very_...".into()), // 5 chars kept + "..." + ], + vec![ + Some("user2".into()), + Some("Bob".into()), + Some("short".into()), + ], + ]; + + assert_eq!(actual, expected) + } + + #[test] + fn test_truncate_unicode_multibyte_chars() { + // Each emoji is 4 bytes but 1 char — byte-based truncation would misbehave here + let fixture = Porcelain(vec![vec![ + Some("🦀🦀🦀🦀🦀🦀".into()), // 6 chars, 24 bytes + Some("hi".into()), // 2 chars, under limit + ]]); + let actual = fixture.truncate(0, 5).into_rows(); + let expected = vec![vec![ + Some("🦀🦀🦀🦀🦀...".into()), // 5 chars kept + "..." + Some("hi".into()), + ]]; + assert_eq!(actual, expected) + } + + #[test] + fn test_truncate_unicode_exactly_at_max_len() { + // String is exactly max_len chars — should NOT be truncated + let fixture = Porcelain(vec![vec![Some("héllo".into())]]); // 5 chars, 6 bytes + let actual = fixture.truncate(0, 5).into_rows(); + let expected = vec![vec![Some("héllo".into())]]; + assert_eq!(actual, expected) + } + + #[test] + fn test_truncate_unicode_exceeds_max_len() { + // 'é' is 2 bytes but 1 char — byte-based slicing would panic or cut wrong + let fixture = Porcelain(vec![vec![Some("héllo world".into())]]); // 11 chars + let actual = fixture.truncate(0, 8).into_rows(); + let expected = vec![vec![Some("héllo wo...".into())]]; // 8 chars kept + "..." + assert_eq!(actual, expected) + } + + #[test] + fn test_swap_cols() { + let info = Porcelain(vec![ + vec![ + Some("user1".into()), + Some("Alice".into()), + Some("30".into()), + ], + vec![Some("user2".into()), Some("Bob".into()), Some("25".into())], + ]); + + let actual = info.swap_cols(0, 1).into_rows(); + + let expected = vec![ + vec![ + Some("Alice".into()), + Some("user1".into()), + Some("30".into()), + ], + vec![Some("Bob".into()), Some("user2".into()), Some("25".into())], + ]; + + assert_eq!(actual, expected) + } + + #[test] + fn test_into_long() { + let info = Info::new() + .add_title("env") + .add_key_value("version", "0.1.0") + .add_key_value("shell", "zsh") + .add_title("conversation") + .add_key_value("id", "000-000-000") + .add_key_value("title", "make agents great again") + .add_title("agent") + .add_key_value("id", "forge") + .add_key_value("model", "sonnet-4"); + + let actual = Porcelain::from(info).into_long(); + let expected = vec![ + vec![ + Some("env".into()), + Some("version".into()), + Some("0.1.0".into()), + ], + vec![Some("env".into()), Some("shell".into()), Some("zsh".into())], + vec![ + Some("conversation".into()), + Some("id".into()), + Some("000-000-000".into()), + ], + vec![ + Some("conversation".into()), + Some("title".into()), + Some("make agents great again".into()), + ], + vec![ + Some("agent".into()), + Some("id".into()), + Some("forge".into()), + ], + vec![ + Some("agent".into()), + Some("model".into()), + Some("sonnet-4".into()), + ], + ]; + + assert_eq!(actual.into_body(), expected) + } + + #[test] + fn test_from_info_single_col() { + let info = Info::new() + .add_title("T1") + .add_value("a1") + .add_value("b1") + .add_title("T2") + .add_value("a2") + .add_value("b2") + .add_title("T3") + .add_value("a3") + .add_value("b3"); + + let actual = Porcelain::from(info).into_rows(); + + let expected = vec![ + // + vec![ + Some(headers::ID.into()), + Some(format!("{}_1", headers::VALUE)), + ], + vec![Some("T1".into()), Some("a1, b1".into())], + vec![Some("T2".into()), Some("a2, b2".into())], + vec![Some("T3".into()), Some("a3, b3".into())], + ]; + + assert_eq!(actual, expected) + } + + #[test] + fn test_into_long_single_col() { + let info = Info::new() + .add_title("T1") + .add_value("a1") + .add_value("b1") + .add_title("T2") + .add_value("a2") + .add_value("b2") + .add_title("T3") + .add_value("a3") + .add_value("b3"); + + let actual = Porcelain::from(info).into_long().into_rows(); + + let expected = vec![ + vec![ + Some(headers::ID.into()), + Some(headers::FIELD.into()), + Some(headers::VALUE.into()), + ], + vec![ + Some("T1".into()), + Some(format!("{}_1", headers::VALUE)), + Some("a1, b1".into()), + ], + vec![ + Some("T2".into()), + Some(format!("{}_1", headers::VALUE)), + Some("a2, b2".into()), + ], + vec![ + Some("T3".into()), + Some(format!("{}_1", headers::VALUE)), + Some("a3, b3".into()), + ], + ]; + + assert_eq!(actual, expected) + } + + #[test] + fn test_display_simple() { + let info = Porcelain(vec![ + vec![ + Some(headers::ID.into()), + Some("name".into()), + Some("age".into()), + ], + vec![ + Some("user1".into()), + Some("Alice".into()), + Some("30".into()), + ], + vec![Some("user2".into()), Some("Bob".into()), Some("25".into())], + ]); + + let actual = info.to_string(); + let expected = [ + // + "ID name age", + "user1 Alice 30", + "user2 Bob 25", + ] + .join("\n"); + + assert_eq!(actual, expected) + } + + #[test] + fn test_display_with_none() { + let info = Porcelain(vec![ + vec![ + Some(headers::ID.into()), + Some("name".into()), + Some("age".into()), + ], + vec![ + Some("user1".into()), + Some("Alice".into()), + Some("30".into()), + ], + vec![Some("user2".into()), None, Some("25".into())], + ]); + + let actual = info.to_string(); + let expected = [ + // + "ID name age", + "user1 Alice 30", + "user2 25", + ] + .join("\n"); + + assert_eq!(actual, expected) + } + + #[test] + fn test_to_case() { + use convert_case::Case; + + let info = Porcelain(vec![ + vec![ + Some(headers::ID.into()), + Some("user_name".into()), + Some("user_age".into()), + ], + vec![ + Some("user1".into()), + Some("Alice Smith".into()), + Some("30".into()), + ], + vec![ + Some("user2".into()), + Some("Bob Jones".into()), + Some("25".into()), + ], + ]); + + let actual = info.to_case(&[1], Case::Snake).into_rows(); + + let expected = vec![ + vec![ + Some(headers::ID.into()), + Some("user_name".into()), + Some("user_age".into()), + ], + vec![ + Some("user1".into()), + Some("alice_smith".into()), + Some("30".into()), + ], + vec![ + Some("user2".into()), + Some("bob_jones".into()), + Some("25".into()), + ], + ]; + + assert_eq!(actual, expected) + } + + #[test] + fn test_to_case_multiple_cols() { + use convert_case::Case; + + let info = Porcelain(vec![ + vec![ + Some("FirstName".into()), + Some("LastName".into()), + Some("Age".into()), + ], + vec![ + Some("Alice".into()), + Some("Smith".into()), + Some("30".into()), + ], + ]); + + let actual = info.to_case(&[0, 1, 2], Case::Kebab).into_rows(); + + let expected = vec![ + vec![ + Some("first-name".into()), + Some("last-name".into()), + Some("age".into()), + ], + vec![ + Some("alice".into()), + Some("smith".into()), + Some("30".into()), + ], + ]; + + assert_eq!(actual, expected) + } + + #[test] + fn test_display_empty() { + let info = Porcelain::new(); + + let actual = info.to_string(); + let expected = ""; + + assert_eq!(actual, expected) + } + + #[test] + fn test_sort_by_single_col() { + let fixture = Porcelain(vec![ + vec![ + Some(headers::ID.into()), + Some("name".into()), + Some("age".into()), + ], + vec![ + Some("user3".into()), + Some("Charlie".into()), + Some("35".into()), + ], + vec![ + Some("user1".into()), + Some("Alice".into()), + Some("30".into()), + ], + vec![Some("user2".into()), Some("Bob".into()), Some("25".into())], + ]); + + let actual = fixture.sort_by(&[1]).into_rows(); + + let expected = vec![ + vec![ + Some(headers::ID.into()), + Some("name".into()), + Some("age".into()), + ], + vec![ + Some("user1".into()), + Some("Alice".into()), + Some("30".into()), + ], + vec![Some("user2".into()), Some("Bob".into()), Some("25".into())], + vec![ + Some("user3".into()), + Some("Charlie".into()), + Some("35".into()), + ], + ]; + + assert_eq!(actual, expected) + } + + #[test] + fn test_sort_by_multiple_cols() { + let fixture = Porcelain(vec![ + vec![ + Some(headers::ID.into()), + Some("city".into()), + Some("name".into()), + ], + vec![ + Some("user3".into()), + Some("NYC".into()), + Some("Charlie".into()), + ], + vec![ + Some("user1".into()), + Some("NYC".into()), + Some("Alice".into()), + ], + vec![Some("user2".into()), Some("LA".into()), Some("Bob".into())], + vec![Some("user4".into()), Some("NYC".into()), Some("Bob".into())], + ]); + + let actual = fixture.sort_by(&[1, 2]).into_rows(); + + let expected = vec![ + vec![ + Some(headers::ID.into()), + Some("city".into()), + Some("name".into()), + ], + vec![Some("user2".into()), Some("LA".into()), Some("Bob".into())], + vec![ + Some("user1".into()), + Some("NYC".into()), + Some("Alice".into()), + ], + vec![Some("user4".into()), Some("NYC".into()), Some("Bob".into())], + vec![ + Some("user3".into()), + Some("NYC".into()), + Some("Charlie".into()), + ], + ]; + + assert_eq!(actual, expected) + } + + #[test] + fn test_sort_by_with_none_values() { + let fixture = Porcelain(vec![ + vec![ + Some(headers::ID.into()), + Some("name".into()), + Some("age".into()), + ], + vec![Some("user3".into()), None, Some("35".into())], + vec![ + Some("user1".into()), + Some("Alice".into()), + Some("30".into()), + ], + vec![Some("user2".into()), Some("Bob".into()), None], + ]); + + let actual = fixture.sort_by(&[1]).into_rows(); + + let expected = vec![ + vec![ + Some(headers::ID.into()), + Some("name".into()), + Some("age".into()), + ], + vec![ + Some("user1".into()), + Some("Alice".into()), + Some("30".into()), + ], + vec![Some("user2".into()), Some("Bob".into()), None], + vec![Some("user3".into()), None, Some("35".into())], + ]; + + assert_eq!(actual, expected) + } + + #[test] + fn test_sort_by_empty_cols() { + let fixture = Porcelain(vec![ + vec![Some(headers::ID.into()), Some("name".into())], + vec![Some("user2".into()), Some("Bob".into())], + vec![Some("user1".into()), Some("Alice".into())], + ]); + + let actual = fixture.sort_by(&[]).into_rows(); + + let expected = vec![ + vec![Some(headers::ID.into()), Some("name".into())], + vec![Some("user2".into()), Some("Bob".into())], + vec![Some("user1".into()), Some("Alice".into())], + ]; + + assert_eq!(actual, expected) + } + + #[test] + fn test_sort_by_empty_porcelain() { + let fixture = Porcelain::new(); + let actual = fixture.sort_by(&[0, 1]).into_rows(); + let expected: Vec>> = vec![]; + + assert_eq!(actual, expected) + } + + #[test] + fn test_sort_by_preserves_header() { + let fixture = Porcelain(vec![ + vec![Some("HEADER".into()), Some("COL1".into())], + vec![Some("z".into()), Some("value".into())], + vec![Some("a".into()), Some("value".into())], + ]); + + let actual = fixture.sort_by(&[0]).into_rows(); + + let expected = vec![ + vec![Some("HEADER".into()), Some("COL1".into())], + vec![Some("a".into()), Some("value".into())], + vec![Some("z".into()), Some("value".into())], + ]; + + assert_eq!(actual, expected) + } + + #[test] + fn test_set_headers_basic() { + let fixture = Porcelain(vec![ + vec![Some("old1".into()), Some("old2".into())], + vec![Some("data1".into()), Some("data2".into())], + ]); + + let actual = fixture.set_headers(vec!["new1", "new2"]).into_rows(); + + let expected = vec![ + vec![Some("new1".into()), Some("new2".into())], + vec![Some("data1".into()), Some("data2".into())], + ]; + + assert_eq!(actual, expected) + } + + #[test] + fn test_set_headers_empty_porcelain() { + let fixture = Porcelain::new(); + + let actual = fixture.set_headers(vec!["header1", "header2"]).into_rows(); + + let expected = vec![vec![Some("header1".into()), Some("header2".into())]]; + + assert_eq!(actual, expected) + } + + #[test] + fn test_set_headers_different_count() { + let fixture = Porcelain(vec![ + vec![ + Some("old1".into()), + Some("old2".into()), + Some("old3".into()), + ], + vec![ + Some("data1".into()), + Some("data2".into()), + Some("data3".into()), + ], + ]); + + let actual = fixture.set_headers(vec!["new1", "new2"]).into_rows(); + + let expected = vec![ + vec![Some("new1".into()), Some("new2".into())], + vec![ + Some("data1".into()), + Some("data2".into()), + Some("data3".into()), + ], + ]; + + assert_eq!(actual, expected) + } + + #[test] + fn test_set_headers_empty_headers() { + let fixture = Porcelain(vec![ + vec![Some("old1".into()), Some("old2".into())], + vec![Some("data1".into()), Some("data2".into())], + ]); + + let actual = fixture.set_headers(Vec::::new()).into_rows(); + + let expected = vec![ + vec![Some("old1".into()), Some("old2".into())], + vec![Some("data1".into()), Some("data2".into())], + ]; + + assert_eq!(actual, expected) + } + + #[test] + fn test_set_headers_with_string_refs() { + let fixture = Porcelain(vec![vec![Some("old".into())], vec![Some("data".into())]]); + + let actual = fixture.set_headers(["ID", "Name", "Age"]).into_rows(); + + let expected = vec![ + vec![Some("ID".into()), Some("Name".into()), Some("Age".into())], + vec![Some("data".into())], + ]; + + assert_eq!(actual, expected) + } +} diff --git a/crates/forge_main/src/vscode.rs b/crates/forge_main/src/vscode.rs new file mode 100644 index 0000000000000000000000000000000000000000..01ee59468afbc2c4f81c98d2086e42050cd758e0 --- /dev/null +++ b/crates/forge_main/src/vscode.rs @@ -0,0 +1,148 @@ +//! VS Code terminal detection and automatic extension installation + +use std::process::Command; + +/// Checks if running in VS Code terminal +/// +/// Detects VS Code by checking for environment variables that are set +/// when running in the VS Code integrated terminal. +pub fn is_vscode_terminal() -> bool { + std::env::var("TERM_PROGRAM") + .map(|val| val == "vscode") + .unwrap_or(false) + || std::env::var("VSCODE_PID").is_ok() + || std::env::var("VSCODE_GIT_ASKPASS_NODE").is_ok() + || std::env::var("VSCODE_GIT_IPC_HANDLE").is_ok() +} + +/// Checks if the Forge VS Code extension is installed +/// +/// Checks VS Code's extension list to see if ForgeCode.forge-vscode is +/// installed. +pub fn is_extension_installed() -> bool { + // Try to list installed extensions + if let Ok(output) = Command::new("code").arg("--list-extensions").output() + && output.status.success() + && let Ok(extensions) = String::from_utf8(output.stdout) + { + return extensions + .lines() + .any(|line| line.trim() == "ForgeCode.forge-vscode"); + } + false +} + +/// Attempts to install the Forge VS Code extension silently +/// +/// Returns Ok(true) if installation was successful, Ok(false) if it failed, +/// or Err if the command couldn't be executed. +pub fn install_extension() -> Result { + let output = Command::new("code") + .arg("--install-extension") + .arg("ForgeCode.forge-vscode") + .arg("--force") + .output()?; + + Ok(output.status.success()) +} + +/// Returns true if we should install the extension +/// +/// This will return true only when: +/// - Not running on native Windows +/// - Running in VS Code terminal +/// - Extension is not installed +pub fn should_install_extension() -> bool { + if cfg!(windows) { + return false; + } + + is_vscode_terminal() && !is_extension_installed() +} + +#[cfg(test)] +mod tests { + use std::env; + use std::ffi::OsString; + use std::sync::{LazyLock, Mutex, MutexGuard}; + + use super::*; + + static ENV_MUTEX: LazyLock> = LazyLock::new(|| Mutex::new(())); + + struct EnvGuard { + key: &'static str, + original_value: Option, + _lock: MutexGuard<'static, ()>, + } + + impl EnvGuard { + fn set(key: &'static str, value: &str) -> Self { + let lock = ENV_MUTEX.lock().unwrap_or_else(|error| error.into_inner()); + let original_value = env::var_os(key); + unsafe { + env::set_var(key, value); + } + Self { key, original_value, _lock: lock } + } + } + + impl Drop for EnvGuard { + fn drop(&mut self) { + match &self.original_value { + Some(value) => unsafe { + env::set_var(self.key, value); + }, + None => unsafe { + env::remove_var(self.key); + }, + } + } + } + + fn with_env_var(key: &'static str, value: &str, test: F) + where + F: FnOnce(), + { + let _guard = EnvGuard::set(key, value); + test(); + } + + #[test] + fn test_is_vscode_terminal_with_term_program() { + with_env_var("TERM_PROGRAM", "vscode", || { + assert!(is_vscode_terminal()); + }); + } + + #[test] + fn test_is_vscode_terminal_with_vscode_pid() { + with_env_var("VSCODE_PID", "12345", || { + assert!(is_vscode_terminal()); + }); + } + + #[test] + fn test_is_vscode_terminal_with_git_askpass() { + with_env_var("VSCODE_GIT_ASKPASS_NODE", "/path/to/node", || { + assert!(is_vscode_terminal()); + }); + } + + #[test] + fn test_is_vscode_terminal_with_git_ipc() { + with_env_var("VSCODE_GIT_IPC_HANDLE", "handle", || { + assert!(is_vscode_terminal()); + }); + } + + #[test] + fn test_should_install_when_in_vscode() { + with_env_var("TERM_PROGRAM", "vscode", || { + // We can't reliably test the actual installation check since it depends + // on the actual VS Code installation, but we can verify the logic + // when in VS Code terminal + assert!(is_vscode_terminal()); + }); + } +} diff --git a/crates/forge_markdown_stream/src/code.rs b/crates/forge_markdown_stream/src/code.rs new file mode 100644 index 0000000000000000000000000000000000000000..0a7603b1ae100256259ca8a32e6ffa4f1b5b2740 --- /dev/null +++ b/crates/forge_markdown_stream/src/code.rs @@ -0,0 +1,124 @@ +//! Code block rendering with syntax highlighting and line wrapping. + +use streamdown_render::code::code_wrap; +use syntect::easy::HighlightLines; +use syntect::highlighting::ThemeSet; +use syntect::parsing::SyntaxSet; +use syntect::util::as_24_bit_terminal_escaped; + +use crate::utils::{ThemeMode, detect_theme_mode}; + +const RESET: &str = "\x1b[0m"; + +/// Code block highlighter using syntect. +pub struct CodeHighlighter { + syntax_set: SyntaxSet, + theme_set: ThemeSet, + theme_mode: ThemeMode, +} + +impl Default for CodeHighlighter { + fn default() -> Self { + Self { + syntax_set: SyntaxSet::load_defaults_newlines(), + theme_set: ThemeSet::load_defaults(), + theme_mode: detect_theme_mode(), + } + } +} + +impl CodeHighlighter { + /// Highlight a single line of code. + fn highlight_line(&self, line: &str, language: Option<&str>) -> String { + let syntax = language + .and_then(|lang| self.syntax_set.find_syntax_by_token(lang)) + .unwrap_or_else(|| self.syntax_set.find_syntax_plain_text()); + + let theme_name = match self.theme_mode { + ThemeMode::Dark => "base16-ocean.dark", + ThemeMode::Light => "InspiredGitHub", + }; + let theme = self.theme_set.themes.get(theme_name).unwrap_or_else(|| { + // Fallback to base16-ocean.dark if theme not found + self.theme_set + .themes + .get("base16-ocean.dark") + .expect("Default theme should exist") + }); + let mut highlighter = HighlightLines::new(syntax, theme); + + match highlighter.highlight_line(line, &self.syntax_set) { + Ok(ranges) => as_24_bit_terminal_escaped(&ranges[..], false), + Err(_) => line.to_string(), + } + } + + /// Render a code line with margin, wrapping if needed. + /// + /// Returns multiple lines if the code exceeds the available width. + pub fn render_code_line( + &self, + line: &str, + language: Option<&str>, + margin: &str, + width: usize, + ) -> Vec { + // Use code_wrap with pretty_broken=true for line wrapping + let (indent, wrapped_lines) = code_wrap(line, width, true); + + let mut result = Vec::new(); + + for (i, code_line) in wrapped_lines.iter().enumerate() { + let highlighted = self.highlight_line(code_line, language); + + // Add continuation indent for wrapped lines + let line_indent = if i == 0 { + "" + } else { + &" ".repeat(indent.min(4) / 2 + 1) + }; + + result.push(format!("{}{}{}{}", margin, line_indent, highlighted, RESET)); + } + + if result.is_empty() { + result.push(format!("{}{}", margin, RESET)); + } + + result + } +} + +#[cfg(test)] +mod tests { + use streamdown_render::code::code_wrap; + + #[test] + fn test_code_wrap_short_line() { + let (indent, lines) = code_wrap("let x = 1;", 80, true); + assert_eq!(indent, 0); + assert_eq!(lines.len(), 1); + assert_eq!(lines[0], "let x = 1;"); + } + + #[test] + fn test_code_wrap_with_indent() { + let (indent, lines) = code_wrap(" let x = 1;", 80, true); + assert_eq!(indent, 4); + assert_eq!(lines.len(), 1); + } + + #[test] + fn test_code_wrap_long_line() { + let long_line = "x".repeat(100); + let (_, lines) = code_wrap(&long_line, 40, true); + assert!(lines.len() > 1); + } + + #[test] + fn test_code_wrap_empty() { + let (indent, lines) = code_wrap("", 80, true); + assert_eq!(indent, 0); + assert_eq!(lines.len(), 1); + } +} diff --git a/crates/forge_markdown_stream/src/heading.rs b/crates/forge_markdown_stream/src/heading.rs new file mode 100644 index 0000000000000000000000000000000000000000..1d1feaebbba30a28de895676abfe023531ee7e76 --- /dev/null +++ b/crates/forge_markdown_stream/src/heading.rs @@ -0,0 +1,280 @@ +//! Heading rendering with theme-based styling. + +use crate::inline::render_inline_content; +use crate::style::{HeadingStyler, InlineStyler}; +use crate::utils::simple_wrap_preserving_spaces; + +/// Render a heading with appropriate styling. +pub fn render_heading( + level: u8, + content: &str, + width: usize, + margin: &str, + styler: &S, +) -> Vec { + // Create the heading prefix (e.g., "# ", "## ", etc.) + let prefix = "#".repeat(level as usize); + + // For h1, uppercase the content before rendering inline elements + let content_to_render = if level == 1 { + content.to_uppercase() + } else { + content.to_string() + }; + + // First render inline elements (bold, italic, etc.) in the content + let rendered_content = render_inline_content(&content_to_render, styler); + + // Adjust width to account for the prefix (e.g., "# " = 2 chars, "## " = 3 + // chars, etc.) + let prefix_display_width = level as usize + 1; + let content_width = width.saturating_sub(prefix_display_width); + let lines = simple_wrap_preserving_spaces(&rendered_content, content_width); + let mut result = Vec::new(); + + for line in lines { + let formatted = match level { + 1 => { + // H1: Bold, left-aligned, uppercase, with dimmed prefix + format!( + "{}\n{}{} {}", + margin, + margin, + styler.dimmed(&styler.h1(&prefix)), + styler.h1(&line) + ) + } + 2 => { + // H2: Bold, bright color, left-aligned, with dimmed prefix + format!( + "{}\n{}{} {}", + margin, + margin, + styler.dimmed(&styler.h2(&prefix)), + styler.h2(&line) + ) + } + 3 => { + format!( + "{}{} {}", + margin, + styler.dimmed(&styler.h3(&prefix)), + styler.h3(&line) + ) + } + 4 => { + format!( + "{}{} {}", + margin, + styler.dimmed(&styler.h4(&prefix)), + styler.h4(&line) + ) + } + 5 => { + format!( + "{}{} {}", + margin, + styler.dimmed(&styler.h5(&prefix)), + styler.h5(&line) + ) + } + _ => { + format!( + "{}{} {}", + margin, + styler.dimmed(&styler.h6(&prefix)), + styler.h6(&line) + ) + } + }; + result.push(formatted); + } + + result +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::theme::TagStyler; + + fn render(level: u8, content: &str) -> String { + render_heading(level, content, 80, " ", &TagStyler).join("\n") + } + + fn render_with_width(level: u8, content: &str, width: usize) -> String { + render_heading(level, content, width, " ", &TagStyler).join("\n") + } + + fn render_with_margin(level: u8, content: &str, margin: &str) -> String { + render_heading(level, content, 80, margin, &TagStyler).join("\n") + } + + #[test] + fn test_h1_simple() { + insta::assert_snapshot!(render(1, "Hello World"), @" + +

#

HELLO WORLD

+ "); + } + + #[test] + fn test_h2_simple() { + insta::assert_snapshot!(render(2, "Chapter One"), @" + +

##

Chapter One

+ "); + } + + #[test] + fn test_h3_simple() { + insta::assert_snapshot!(render(3, "Section Title"), @"

###

Section Title

"); + } + + #[test] + fn test_h4_simple() { + insta::assert_snapshot!(render(4, "Subsection"), @"

####

Subsection

"); + } + + #[test] + fn test_h5_simple() { + insta::assert_snapshot!(render(5, "Minor Heading"), @"
#####
Minor Heading
"); + } + + #[test] + fn test_h6_simple() { + insta::assert_snapshot!(render(6, "Smallest Heading"), @"
######
Smallest Heading
"); + } + + #[test] + fn test_h1_with_inline_bold() { + insta::assert_snapshot!(render(1, "Hello **bold** world"), @" + +

#

HELLO BOLD WORLD

+ "); + } + + #[test] + fn test_h2_with_inline_italic() { + insta::assert_snapshot!(render(2, "Hello *italic* text"), @" + +

##

Hello italic text

+ "); + } + + #[test] + fn test_h3_with_code() { + insta::assert_snapshot!(render(3, "Using `code` here"), @"

###

Using code here

"); + } + + #[test] + fn test_heading_level_beyond_6() { + // Level 7+ should fall through to h6 styling + insta::assert_snapshot!(render(7, "Level Seven"), @"
#######
Level Seven
"); + insta::assert_snapshot!(render(10, "Level Ten"), @"
##########
Level Ten
"); + } + + #[test] + fn test_empty_content() { + insta::assert_snapshot!(render(1, ""), @" + +

#

+ "); + } + + #[test] + fn test_custom_margin() { + insta::assert_snapshot!(render_with_margin(1, "Title", " "), @" + +

#

TITLE

+ "); + insta::assert_snapshot!(render_with_margin(3, "Section", ">>> "), @">>>

###

Section

"); + } + + #[test] + fn test_no_margin() { + insta::assert_snapshot!(render_with_margin(1, "Title", ""), @" + +

#

TITLE

+ "); + insta::assert_snapshot!(render_with_margin(3, "Section", ""), @"

###

Section

"); + } + + #[test] + fn test_wrapping_narrow_width() { + insta::assert_snapshot!(render_with_width(1, "This is a very long heading that should wrap", 20), @" + +

#

THIS IS A VERY

+ +

#

LONG HEADING THAT

+ +

#

SHOULD WRAP

+ "); + } + + #[test] + fn test_h3_wrapping() { + insta::assert_snapshot!(render_with_width(3, "A long section title that wraps", 15), @" +

###

A long

+

###

section

+

###

title that

+

###

wraps

+ "); + } + + #[test] + fn test_h3_wrapping_preserves_korean_word_spaces() { + let actual = render_with_width(3, "한글 공백 보존 확인", 12); + + insta::assert_snapshot!(actual, @r" +

###

한글

+

###

공백

+

###

보존

+

###

확인

+ "); + } + + #[test] + fn test_h3_wrapping_splits_long_tokens() { + let actual = render_with_width(3, "supercalifragilistic", 12); + + insta::assert_snapshot!(actual, @r" +

###

supercal

+

###

ifragili

+

###

stic

+ "); + } + + #[test] + fn test_special_characters() { + insta::assert_snapshot!(render(2, "Hello & Goodbye < World >"), @" + +

##

Hello & Goodbye < World >

+ "); + } + + #[test] + fn test_heading_with_link() { + insta::assert_snapshot!(render(3, "See [documentation](https://example.com)"), @r#"

###

See documentation

"#); + } + + #[test] + fn test_mixed_inline_styles() { + insta::assert_snapshot!(render(2, "**Bold** and *italic* and `code`"), @" + +

##

Bold and italic and code

+ "); + } + + #[test] + fn test_all_levels_structure() { + // H1 and H2 have extra newline prefix + let h1 = render(1, "H1"); + let h2 = render(2, "H2"); + let h3 = render(3, "H3"); + + assert!(h1.contains("\n"), "H1 should have newline"); + assert!(h2.contains("\n"), "H2 should have newline"); + assert!(!h3.starts_with("\n"), "H3 should not start with newline"); + } +} diff --git a/crates/forge_markdown_stream/src/inline.rs b/crates/forge_markdown_stream/src/inline.rs new file mode 100644 index 0000000000000000000000000000000000000000..06f9b21605361851c2a0a888e5c165fc44c07777 --- /dev/null +++ b/crates/forge_markdown_stream/src/inline.rs @@ -0,0 +1,156 @@ +//! Inline content rendering with theme-based formatting. + +use streamdown_parser::{InlineElement, InlineParser}; + +use crate::style::InlineStyler; + +/// Render inline elements to a string using a styler. +pub fn render_inline_content(content: &str, styler: &S) -> String { + render_inline_elements(&InlineParser::new().parse(content), styler) +} + +/// Render inline elements to a string using a styler. +pub fn render_inline_elements(elements: &[InlineElement], styler: &S) -> String { + let mut result = String::new(); + for element in elements { + match element { + InlineElement::Text(text) => { + result.push_str(&styler.text(text)); + } + InlineElement::Bold(text) => { + result.push_str(&styler.bold(text)); + } + InlineElement::Italic(text) => { + result.push_str(&styler.italic(text)); + } + InlineElement::BoldItalic(text) => { + result.push_str(&styler.bold_italic(text)); + } + InlineElement::Strikeout(text) => { + result.push_str(&styler.strikethrough(text)); + } + InlineElement::Underline(text) => { + result.push_str(&styler.underline(text)); + } + InlineElement::Code(text) => { + result.push_str(&styler.code(text)); + } + InlineElement::Link { text, url } => { + result.push_str(&styler.link(text, url)); + } + InlineElement::Image { alt, url } => { + result.push_str(&styler.image(alt, url)); + } + InlineElement::Footnote(text) => { + result.push_str(&styler.footnote(text)); + } + } + } + + result +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::theme::TagStyler; + + fn render(content: &str) -> String { + render_inline_content(content, &TagStyler) + } + + #[test] + fn test_plain_text() { + insta::assert_snapshot!(render("hello world"), @"hello world"); + } + + #[test] + fn test_html_entities() { + insta::assert_snapshot!(render("& < > ""), @r#"& < > ""#); + } + + #[test] + fn test_bold() { + insta::assert_snapshot!(render("**bold**"), @"bold"); + } + + #[test] + fn test_italic() { + insta::assert_snapshot!(render("*italic*"), @"italic"); + } + + #[test] + fn test_bold_italic() { + insta::assert_snapshot!(render("***text***"), @"text"); + } + + #[test] + fn test_strikethrough() { + insta::assert_snapshot!(render("~~struck~~"), @"struck"); + } + + #[test] + fn test_code() { + insta::assert_snapshot!(render("`code`"), @"code"); + } + + #[test] + fn test_underline() { + insta::assert_snapshot!(render("__underline__"), @"underline"); + } + + #[test] + fn test_link() { + insta::assert_snapshot!(render("[click](https://example.com)"), @r#"click"#); + } + + #[test] + fn test_image() { + insta::assert_snapshot!(render("![alt](image.png)"), @r#"alt"#); + } + + #[test] + fn test_mixed() { + insta::assert_snapshot!(render("hello **bold** and *italic*"), @"hello bold and italic"); + } + + #[test] + fn test_multiple_bold() { + insta::assert_snapshot!(render("**one** and **two**"), @"one and two"); + } + + #[test] + fn test_entities_in_bold() { + insta::assert_snapshot!(render("**&**"), @"&"); + } + + #[test] + fn test_entities_in_link() { + insta::assert_snapshot!(render("[<click>](https://example.com)"), @r#""#); + } + + #[test] + fn test_code_content() { + insta::assert_snapshot!(render("`let x = 1;`"), @"let x = 1;"); + } + + #[test] + fn test_link_special_url() { + insta::assert_snapshot!(render("[link](https://example.com/path?q=1&b=2)"), @r#"link"#); + } + + #[test] + fn test_empty() { + insta::assert_snapshot!(render(""), @""); + } + + #[test] + fn test_whitespace() { + insta::assert_snapshot!(render("hello world"), @"hello world"); + } + + #[test] + fn test_image_empty_alt() { + insta::assert_snapshot!(render("![](image.png)"), @r#""#); + } +} diff --git a/crates/forge_markdown_stream/src/lib.rs b/crates/forge_markdown_stream/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..8f8e3d1c57cafffd59edea913caee5fb6bbd5f7b --- /dev/null +++ b/crates/forge_markdown_stream/src/lib.rs @@ -0,0 +1,233 @@ +//! Forge Markdown Stream - Streaming markdown renderer for terminal output. +//! +//! This crate provides a streaming markdown renderer optimized for LLM output. +//! It renders markdown with syntax highlighting, styled headings, tables, +//! lists, and more. +//! +//! # Example +//! +//! ```no_run +//! use forge_markdown_stream::StreamdownRenderer; +//! use std::io; +//! +//! fn main() -> io::Result<()> { +//! let mut renderer = StreamdownRenderer::new(io::stdout(), 80); +//! +//! // Push tokens as they arrive from LLM +//! renderer.push("Hello ")?; +//! renderer.push("**world**!\n")?; +//! +//! // Finish rendering +//! let _ = renderer.finish()?; +//! Ok(()) +//! } +//! ``` + +mod code; +mod heading; +mod inline; +mod list; +mod renderer; +mod repair; +mod style; +mod table; +mod theme; +mod utils; + +use std::io::{self, Write}; + +pub use renderer::Renderer; +pub use repair::repair_line; +pub use streamdown_parser::Parser; +pub use theme::{Style, Theme}; + +/// Streaming markdown renderer for terminal output. +/// +/// Buffers incoming tokens and renders complete lines with syntax highlighting, +/// styled headings, tables, lists, and more. +/// +/// The renderer is generic over the writer type `W`, which must implement +/// `Write`. +pub struct StreamdownRenderer { + parser: Parser, + renderer: Renderer, + line_buffer: String, +} + +impl StreamdownRenderer { + /// Create a new renderer with the given writer and terminal width. + pub fn new(writer: W, width: usize) -> Self { + Self { + parser: Parser::new(), + renderer: Renderer::new(writer, width), + line_buffer: String::new(), + } + } + + /// Create a new renderer with a custom theme. + pub fn with_theme(writer: W, width: usize, theme: Theme) -> Self { + Self { + parser: Parser::new(), + renderer: Renderer::with_theme(writer, width, theme), + line_buffer: String::new(), + } + } + + /// Push a token to the renderer. + /// + /// Tokens are buffered until a complete line is received, then rendered. + pub fn push(&mut self, token: &str) -> io::Result<()> { + self.line_buffer.push_str(token); + + while let Some(pos) = self.line_buffer.find('\n') { + let line = self.line_buffer.get(..pos).unwrap_or("").to_string(); + + for repaired in repair_line(&line, self.parser.state()) { + for event in self.parser.parse_line(&repaired) { + self.renderer.render_event(&event)?; + } + } + + self.line_buffer = self.line_buffer.get(pos + 1..).unwrap_or("").to_string(); + } + Ok(()) + } + + /// Finish rendering, flushing any remaining buffered content. + /// Returns the underlying writer. + pub fn finish(mut self) -> io::Result<()> { + if !self.line_buffer.is_empty() { + for repaired in repair_line(&self.line_buffer, self.parser.state()) { + for event in self.parser.parse_line(&repaired) { + self.renderer.render_event(&event)?; + } + } + } + for event in self.parser.finalize() { + self.renderer.render_event(&event)?; + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::StreamdownRenderer; + + fn fixture_rendered_output(markdown: &str, width: usize) -> String { + let mut output = Vec::new(); + let mut fixture = StreamdownRenderer::new(&mut output, width); + fixture.push(markdown).unwrap(); + fixture.finish().unwrap(); + + let actual = strip_ansi_escapes::strip(output); + String::from_utf8(actual) + .unwrap() + .trim_matches('\n') + .to_string() + } + + fn fixture_rendered_output_from_chunks(chunks: &[&str], width: usize) -> String { + let mut output = Vec::new(); + let mut fixture = StreamdownRenderer::new(&mut output, width); + for chunk in chunks { + fixture.push(chunk).unwrap(); + } + fixture.finish().unwrap(); + + let actual = strip_ansi_escapes::strip(output); + String::from_utf8(actual) + .unwrap() + .trim_matches('\n') + .to_string() + } + + #[test] + fn test_streaming_renderer_preserves_korean_spacing_in_structured_markdown() { + let fixture = concat!( + "## 구현 요약\n", + "- 각 서비스에서 metadata key를 개별 수정하지 않고, object storage 공통 레이어에서 일괄 정규화하도록 반영했습니다.\n", + "## 검토 사항\n", + "- 본 수정은 업로드 시 metadata header 이름 문제를 해결합니다.\n", + "- 추가적인 권한 정책, bucket policy, reverse proxy 제한이 있으면 별도 오류가 발생할 수 있습니다.\n", + ); + let actual = fixture_rendered_output(fixture, 200); + let expected = concat!( + "## 구현 요약\n", + "• 각 서비스에서 metadata key를 개별 수정하지 않고, object storage 공통 레이어에서 일괄 정규화하도록 반영했습니다.\n", + "\n", + "## 검토 사항\n", + "• 본 수정은 업로드 시 metadata header 이름 문제를 해결합니다.\n", + "• 추가적인 권한 정책, bucket policy, reverse proxy 제한이 있으면 별도 오류가 발생할 수 있습니다.", + ); + + assert_eq!(actual, expected); + } + + #[test] + fn test_streaming_renderer_preserves_korean_spacing_when_structured_tail_arrives_in_chunks() { + let fixture = [ + "## 검토 결과\n", + "- 본 사례는 스트리밍 마크다운 렌더링의 공백 재조합 문제와 관련이 있습니다.\n", + "- 핵심 구현은 공백 보존 래퍼에 위치합니다.\n", + "- 회귀 테스트는 스트리밍 렌더러 검증 항목에 추가되어 있습니다.\n\n", + "후속 작업은 다음과 같습니다.\n", + "1. 변경 사항을 검토 가능한 형식으로 정리합니다.\n", + "2. 실제 대화 출력과 유사한 통합 테스트 범위를 ", + "확장합니다.", + ]; + let actual = fixture_rendered_output_from_chunks(&fixture, 200); + let expected = concat!( + "## 검토 결과\n", + "• 본 사례는 스트리밍 마크다운 렌더링의 공백 재조합 문제와 관련이 있습니다.\n", + "• 핵심 구현은 공백 보존 래퍼에 위치합니다.\n", + "• 회귀 테스트는 스트리밍 렌더러 검증 항목에 추가되어 있습니다.\n", + "\n", + "후속 작업은 다음과 같습니다.\n", + "1. 변경 사항을 검토 가능한 형식으로 정리합니다.\n", + "2. 실제 대화 출력과 유사한 통합 테스트 범위를 확장합니다.", + ); + + assert_eq!(actual, expected); + } + + #[test] + fn test_streaming_renderer_wraps_blockquotes_with_prefix_width_and_long_tokens() { + let fixture = "> supercalifragilistic\n> 한글 공백\n"; + let actual = fixture_rendered_output(fixture, 10); + let expected = concat!( + "│ supercal\n", + "│ ifragili\n", + "│ stic\n", + "│ 한글\n", + "│ 공백" + ); + + assert_eq!(actual, expected); + } + + #[test] + fn test_streaming_renderer_wraps_blockquote_links_without_losing_separator() { + let fixture = "> [링크](https://example.com/very/long/path) 설명\n"; + let actual = fixture_rendered_output(fixture, 20); + let expected = concat!( + "│ 링크\n", + "│ (https://example.c\n", + "│ om/very/long/path)\n", + "│ 설명" + ); + + assert_eq!(actual, expected); + } + + #[test] + fn test_streaming_renderer_wraps_nested_blockquotes_with_correct_prefix_width() { + let fixture = ">> supercalifragilistic\n"; + let actual = fixture_rendered_output(fixture, 12); + let expected = concat!("│ │ supercal\n", "│ │ ifragili\n", "│ │ stic"); + + assert_eq!(actual, expected); + } +} diff --git a/crates/forge_markdown_stream/src/list.rs b/crates/forge_markdown_stream/src/list.rs new file mode 100644 index 0000000000000000000000000000000000000000..836b650f780d1e81d8b8347d45c0fb3acc559387 --- /dev/null +++ b/crates/forge_markdown_stream/src/list.rs @@ -0,0 +1,687 @@ +//! List rendering with nested indentation and bullet cycling. + +use streamdown_ansi::utils::visible_length; +use streamdown_parser::ListBullet; + +use crate::inline::render_inline_content; +use crate::style::{InlineStyler, ListStyler}; +use crate::utils::wrap_text_preserving_spaces; + +/// Bullet characters for dash lists at different nesting levels. +const BULLETS_DASH: [&str; 4] = ["•", "◦", "▪", "‣"]; + +/// Bullet characters for asterisk lists at different nesting levels. +const BULLETS_ASTERISK: [&str; 4] = ["∗", "⁎", "✱", "✳"]; + +/// Bullet characters for plus lists at different nesting levels. +const BULLETS_PLUS: [&str; 4] = ["⊕", "⊙", "⊛", "⊜"]; + +/// Checkbox characters for task list items. +const CHECKBOX_UNCHECKED: &str = ""; +const CHECKBOX_CHECKED: &str = ""; + +/// Strips checkbox prefix from content and returns (checkbox_char, +/// remaining_content). Returns None if no checkbox is found at the start. +fn strip_checkbox_prefix(content: &str) -> Option<(&'static str, &str)> { + if let Some(rest) = content.strip_prefix("[ ] ") { + Some((CHECKBOX_UNCHECKED, rest)) + } else if let Some(rest) = content + .strip_prefix("[x] ") + .or_else(|| content.strip_prefix("[X] ")) + { + Some((CHECKBOX_CHECKED, rest)) + } else if content == "[ ]" { + Some((CHECKBOX_UNCHECKED, "")) + } else if content == "[x]" || content == "[X]" { + Some((CHECKBOX_CHECKED, "")) + } else { + None + } +} + +/// List rendering state for tracking nesting and numbering. +#[derive(Default)] +pub struct ListState { + /// Stack of (indent, is_ordered) for nested lists + stack: Vec<(usize, bool)>, + /// Current ordered list numbers at each level + numbers: Vec, + /// Whether we're in a "pending" state (saw ListEnd but might continue) + pending_reset: bool, +} + +impl ListState { + pub fn level(&self) -> usize { + self.stack.len() + } + + pub fn push(&mut self, indent: usize, ordered: bool) { + self.stack.push((indent, ordered)); + self.numbers.push(0); + } + + pub fn pop(&mut self) { + self.stack.pop(); + self.numbers.pop(); + } + + pub fn next_number(&mut self) -> usize { + if let Some(n) = self.numbers.last_mut() { + *n += 1; + *n + } else { + 1 + } + } + + pub fn adjust_for_indent(&mut self, indent: usize, ordered: bool) { + // Pop levels that are deeper than current + while let Some((stack_indent, _)) = self.stack.last() { + if *stack_indent > indent { + self.pop(); + } else { + break; + } + } + + // Check if we need to push a new level + let need_push = self.stack.last().map(|(i, _)| indent > *i).unwrap_or(true); + if need_push { + self.push(indent, ordered); + } + } + + pub fn reset(&mut self) { + self.stack.clear(); + self.numbers.clear(); + self.pending_reset = false; + } + + /// Mark list as pending reset (saw ListEnd, but might continue with more + /// items) + pub fn mark_pending_reset(&mut self) { + self.pending_reset = true; + } + + /// Resume list if it was pending reset (new list item arrived) + fn resume_if_pending(&mut self) { + self.pending_reset = false; + } +} + +/// Render a list item. +pub fn render_list_item( + indent: usize, + bullet: &ListBullet, + content: &str, + width: usize, + margin: &str, + styler: &S, + list_state: &mut ListState, +) -> Vec { + // Resume list if it was pending reset (continues after empty line) + list_state.resume_if_pending(); + + // Adjust list state for current indent + let ordered = matches!(bullet, ListBullet::Ordered(_)); + list_state.adjust_for_indent(indent, ordered); + + let level = list_state.level().saturating_sub(1); + + // Check for checkbox at start of content + let (checkbox_prefix, actual_content) = match strip_checkbox_prefix(content) { + Some((checkbox, rest)) => { + let styled = if checkbox == CHECKBOX_CHECKED { + styler.checkbox_checked(checkbox) + } else { + styler.checkbox_unchecked(checkbox) + }; + (format!("{} ", styled), rest) + } + None => (String::new(), content), + }; + + // Calculate marker - use our own counter for ordered lists to work around + // the parser bug that normalizes all numbers to 1 + let marker = match bullet { + ListBullet::Ordered(_) => { + let num = list_state.next_number(); + format!("{}.", num) + } + ListBullet::PlusExpand => "⊞".to_string(), + ListBullet::Dash => BULLETS_DASH + .get(level % BULLETS_DASH.len()) + .unwrap_or(&"•") + .to_string(), + ListBullet::Asterisk => BULLETS_ASTERISK + .get(level % BULLETS_ASTERISK.len()) + .unwrap_or(&"∗") + .to_string(), + ListBullet::Plus => BULLETS_PLUS + .get(level % BULLETS_PLUS.len()) + .unwrap_or(&"⊕") + .to_string(), + }; + + // Calculate indentation + let indent_spaces = indent * 2; + let marker_width = visible_length(&marker); + let checkbox_width = if checkbox_prefix.is_empty() { 0 } else { 2 }; // checkbox + space + let content_indent = indent_spaces + marker_width + 1 + checkbox_width; + + // Color the marker based on bullet type + let colored_marker = match bullet { + ListBullet::Ordered(_) => styler.number(&marker), + ListBullet::Dash => styler.bullet_dash(&marker), + ListBullet::Asterisk => styler.bullet_asterisk(&marker), + ListBullet::Plus => styler.bullet_plus(&marker), + ListBullet::PlusExpand => styler.bullet_plus_expand(&marker), + }; + + // Parse and render inline content + let rendered_content = render_inline_content(actual_content, styler); + + // Build prefixes + let first_prefix = format!( + "{}{}{} {}", + margin, + " ".repeat(indent_spaces), + colored_marker, + checkbox_prefix + ); + let next_prefix = format!("{}{}", margin, " ".repeat(content_indent)); + + // The wrapper takes separate first/next content budgets. In list items the + // visible widths of the first and continuation prefixes match by + // construction, but keep the calculation explicit so merge logic stays + // obvious. + const MIN_CONTENT_WIDTH: usize = 5; + let first_content_width = width + .saturating_sub(visible_length(&first_prefix)) + .max(MIN_CONTENT_WIDTH); + let next_content_width = width + .saturating_sub(visible_length(&next_prefix)) + .max(MIN_CONTENT_WIDTH); + + // Wrap the content + let wrapped = wrap_text_preserving_spaces( + &rendered_content, + first_content_width, + next_content_width, + &first_prefix, + &next_prefix, + ); + + if wrapped.is_empty() { + vec![first_prefix] + } else { + wrapped + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::theme::{TagStyler, Theme}; + + fn render(indent: usize, bullet: ListBullet, content: &str) -> String { + let mut state = ListState::default(); + render_list_item(indent, &bullet, content, 80, " ", &TagStyler, &mut state).join("\n") + } + + fn render_with_state( + indent: usize, + bullet: ListBullet, + content: &str, + state: &mut ListState, + ) -> String { + render_list_item(indent, &bullet, content, 80, " ", &TagStyler, state).join("\n") + } + + fn render_with_width(indent: usize, bullet: ListBullet, content: &str, width: usize) -> String { + let mut state = ListState::default(); + render_list_item( + indent, &bullet, content, width, " ", &TagStyler, &mut state, + ) + .join("\n") + } + + fn render_visible_with_width( + indent: usize, + bullet: ListBullet, + content: &str, + width: usize, + ) -> String { + let mut state = ListState::default(); + let actual = render_list_item( + indent, + &bullet, + content, + width, + " ", + &Theme::default(), + &mut state, + ) + .join("\n"); + let stripped = strip_ansi_escapes::strip(actual.as_bytes()); + + String::from_utf8(stripped).unwrap() + } + + #[test] + fn test_unordered_dash() { + insta::assert_snapshot!(render(0, ListBullet::Dash, "Item one"), @" • Item one"); + } + + #[test] + fn test_unordered_asterisk() { + insta::assert_snapshot!(render(0, ListBullet::Asterisk, "Item two"), @" ∗ Item two"); + } + + #[test] + fn test_unordered_plus() { + insta::assert_snapshot!(render(0, ListBullet::Plus, "Item three"), @" ⊕ Item three"); + } + + #[test] + fn test_ordered_item() { + insta::assert_snapshot!(render(0, ListBullet::Ordered(1), "First item"), @" 1. First item"); + } + + #[test] + fn test_ordered_sequential() { + let mut state = ListState::default(); + let first = render_with_state(0, ListBullet::Ordered(1), "First", &mut state); + let second = render_with_state(0, ListBullet::Ordered(1), "Second", &mut state); + let third = render_with_state(0, ListBullet::Ordered(1), "Third", &mut state); + + insta::assert_snapshot!(first, @" 1. First"); + insta::assert_snapshot!(second, @" 2. Second"); + insta::assert_snapshot!(third, @" 3. Third"); + } + + #[test] + fn test_plus_expand() { + insta::assert_snapshot!(render(0, ListBullet::PlusExpand, "Expandable"), @" ⊞ Expandable"); + } + + #[test] + fn test_nested_indent_level_1() { + let mut state = ListState::default(); + // First item at level 0 + let _ = render_with_state(0, ListBullet::Dash, "Parent", &mut state); + // Nested item at indent 1 + let nested = render_with_state(1, ListBullet::Dash, "Child", &mut state); + insta::assert_snapshot!(nested, @" ◦ Child"); + } + + #[test] + fn test_nested_indent_level_2() { + let mut state = ListState::default(); + let _ = render_with_state(0, ListBullet::Dash, "Level 0", &mut state); + let _ = render_with_state(1, ListBullet::Dash, "Level 1", &mut state); + let level2 = render_with_state(2, ListBullet::Dash, "Level 2", &mut state); + insta::assert_snapshot!(level2, @" ▪ Level 2"); + } + + #[test] + fn test_bullet_cycling() { + let mut state = ListState::default(); + let l0 = render_with_state(0, ListBullet::Dash, "L0", &mut state); + let l1 = render_with_state(1, ListBullet::Dash, "L1", &mut state); + let l2 = render_with_state(2, ListBullet::Dash, "L2", &mut state); + let l3 = render_with_state(3, ListBullet::Dash, "L3", &mut state); + let l4 = render_with_state(4, ListBullet::Dash, "L4", &mut state); // cycles back + + assert!(l0.contains("•"), "Level 0 should use •"); + assert!(l1.contains("◦"), "Level 1 should use ◦"); + assert!(l2.contains("▪"), "Level 2 should use ▪"); + assert!(l3.contains("‣"), "Level 3 should use ‣"); + assert!(l4.contains("•"), "Level 4 should cycle back to •"); + } + + #[test] + fn test_inline_bold() { + insta::assert_snapshot!(render(0, ListBullet::Dash, "Item with **bold** text"), @" • Item with bold text"); + } + + #[test] + fn test_inline_italic() { + insta::assert_snapshot!(render(0, ListBullet::Dash, "Item with *italic* text"), @" • Item with italic text"); + } + + #[test] + fn test_inline_code() { + insta::assert_snapshot!(render(0, ListBullet::Dash, "Item with `code` text"), @" • Item with code text"); + } + + #[test] + fn test_inline_link() { + insta::assert_snapshot!(render(0, ListBullet::Dash, "See [link](https://example.com)"), @r#" • See link"#); + } + + #[test] + fn test_empty_content() { + insta::assert_snapshot!(render(0, ListBullet::Dash, ""), @" •"); + } + + #[test] + fn test_wrapping_long_content() { + let result = render_with_width( + 0, + ListBullet::Dash, + "This is a very long list item that should wrap to multiple lines", + 40, + ); + insta::assert_snapshot!(result, @r" + • This is a very long + list item that should wrap to + multiple lines + "); + } + + // Use Theme rather than TagStyler: visible_length skips ANSI codes but + // counts TagStyler's literal markers, which would inflate widths. + + #[test] + fn test_wrapping_respects_width() { + let theme = crate::theme::Theme::dark(); + let mut state = ListState::default(); + let lines = render_list_item( + 0, + &ListBullet::Dash, + "word1 word2 word3 word4 word5 word6 word7 word8", + 20, + " ", + &theme, + &mut state, + ); + for line in &lines { + let vis = visible_length(line); + assert!(vis <= 20, "line {line:?} has visible width {vis} > 20"); + } + } + + #[test] + fn test_wrapping_respects_width_nested() { + let theme = crate::theme::Theme::dark(); + let mut state = ListState::default(); + let _ = render_list_item(0, &ListBullet::Dash, "parent", 30, " ", &theme, &mut state); + let lines = render_list_item( + 1, + &ListBullet::Dash, + "alpha beta gamma delta epsilon zeta eta theta", + 30, + " ", + &theme, + &mut state, + ); + for line in &lines { + let vis = visible_length(line); + assert!(vis <= 30, "line {line:?} has visible width {vis} > 30"); + } + } + + #[test] + fn test_wrapping_preserves_korean_word_spaces() { + let actual = render_with_width(0, ListBullet::Dash, "한글 공백 보존 확인", 8); + let expected = " • 한글\n 공백\n 보존\n 확인"; + + pretty_assertions::assert_eq!(actual, expected); + } + + #[test] + fn test_wrapping_respects_bullet_prefix_width() { + let actual = render_with_width(0, ListBullet::Dash, "한글 공백", 6); + let expected = " • 한글\n 공백"; + + pretty_assertions::assert_eq!(actual, expected); + } + + #[test] + fn test_wrapping_respects_checkbox_prefix_width() { + let actual = render_with_width(0, ListBullet::Dash, "[ ] 한글 공백", 8); + let expected = " • 한글\n 공백"; + + pretty_assertions::assert_eq!(actual, expected); + } + + #[test] + fn test_wrapping_respects_multidigit_ordered_prefix_width() { + let mut fixture = ListState::default(); + for index in 1..10 { + let _ = render_list_item( + 0, + &ListBullet::Ordered(1), + &format!("예시 {index}"), + 8, + " ", + &TagStyler, + &mut fixture, + ); + } + let actual = render_list_item( + 0, + &ListBullet::Ordered(1), + "한글 공백", + 8, + " ", + &TagStyler, + &mut fixture, + ) + .join("\n"); + let expected = " 10. 한글\n 공백"; + + pretty_assertions::assert_eq!(actual, expected); + } + + #[test] + fn test_wrapping_splits_long_tokens() { + let actual = render_with_width(0, ListBullet::Dash, "supercalifragilistic", 10); + let expected = " • super\n califr\n agilis\n tic"; + + pretty_assertions::assert_eq!(actual, expected); + } + + #[test] + fn test_wrapping_preserves_link_breaks() { + let actual = render_visible_with_width( + 0, + ListBullet::Dash, + "[링크](https://example.com/very/long/path) 설명", + 14, + ); + let expected = concat!( + " • 링크\n", + " (https://e\n", + " xample.com\n", + " /very/long\n", + " /path)\n", + " 설명" + ); + + pretty_assertions::assert_eq!(actual, expected); + } + + #[test] + fn test_list_state_reset() { + let mut state = ListState::default(); + let _ = render_with_state(0, ListBullet::Ordered(1), "First", &mut state); + let _ = render_with_state(0, ListBullet::Ordered(1), "Second", &mut state); + state.reset(); + let after_reset = render_with_state(0, ListBullet::Ordered(1), "New first", &mut state); + insta::assert_snapshot!(after_reset, @" 1. New first"); + } + + #[test] + fn test_dedent_resets_nested_numbers() { + let mut state = ListState::default(); + let _ = render_with_state(0, ListBullet::Ordered(1), "Parent 1", &mut state); + let _ = render_with_state(1, ListBullet::Ordered(1), "Child 1", &mut state); + let _ = render_with_state(1, ListBullet::Ordered(1), "Child 2", &mut state); + let parent2 = render_with_state(0, ListBullet::Ordered(1), "Parent 2", &mut state); + insta::assert_snapshot!(parent2, @" 2. Parent 2"); + } + + #[test] + fn test_mixed_ordered_unordered() { + let mut state = ListState::default(); + let ordered = render_with_state(0, ListBullet::Ordered(1), "Ordered item", &mut state); + state.reset(); + let mut state2 = ListState::default(); + let unordered = render_with_state(0, ListBullet::Dash, "Unordered item", &mut state2); + + assert!(ordered.contains(""), "Ordered should use number style"); + assert!( + unordered.contains(""), + "Unordered should use dash style" + ); + } + + #[test] + fn test_list_state_level() { + let mut state = ListState::default(); + assert_eq!(state.level(), 0); + + state.push(0, false); + assert_eq!(state.level(), 1); + + state.push(1, false); + assert_eq!(state.level(), 2); + + state.pop(); + assert_eq!(state.level(), 1); + } + + mod checkbox { + use super::*; + + mod strip_checkbox_prefix_tests { + use super::*; + + #[test] + fn valid_patterns() { + // (input, expected_checkbox, expected_remaining) + let cases = [ + ("[ ] Task", Some((CHECKBOX_UNCHECKED, "Task"))), + ("[x] Done", Some((CHECKBOX_CHECKED, "Done"))), + ("[X] Done", Some((CHECKBOX_CHECKED, "Done"))), + ("[ ]", Some((CHECKBOX_UNCHECKED, ""))), + ("[x]", Some((CHECKBOX_CHECKED, ""))), + ("[X]", Some((CHECKBOX_CHECKED, ""))), + ]; + + for (input, expected) in cases { + let actual = strip_checkbox_prefix(input); + assert_eq!(actual, expected, "input: {input:?}"); + } + } + + #[test] + fn invalid_patterns() { + let cases = [ + "[] text", // no space inside brackets + "[y] text", // wrong character + "prefix [ ] suffix", // not at start + "array[x]", // array index syntax + "Just plain text", // no brackets + " [ ] Task", // leading whitespace + "[X]Task", // no space after checkbox + ]; + + for input in cases { + let actual = strip_checkbox_prefix(input); + assert_eq!(actual, None, "input: {input:?} should not match"); + } + } + } + + mod render_tests { + use super::*; + + #[test] + fn checkbox_unchecked() { + insta::assert_snapshot!( + render(0, ListBullet::Dash, "[ ] Task to do"), + @" • Task to do" + ); + } + + #[test] + fn checkbox_checked_lowercase() { + insta::assert_snapshot!( + render(0, ListBullet::Dash, "[x] Completed task"), + @" • Completed task" + ); + } + + #[test] + fn checkbox_checked_uppercase() { + insta::assert_snapshot!( + render(0, ListBullet::Dash, "[X] Another completed task"), + @" • Another completed task" + ); + } + + #[test] + fn checkbox_unchecked_empty_content() { + insta::assert_snapshot!( + render(0, ListBullet::Dash, "[ ]"), + @" • " + ); + } + + #[test] + fn checkbox_checked_empty_content() { + insta::assert_snapshot!( + render(0, ListBullet::Dash, "[x]"), + @" • " + ); + } + + #[test] + fn checkbox_with_ordered_list() { + insta::assert_snapshot!( + render(0, ListBullet::Ordered(1), "[ ] Ordered task"), + @" 1. Ordered task" + ); + } + } + + mod no_false_positives { + use super::*; + + #[test] + fn empty_brackets() { + insta::assert_snapshot!( + render(0, ListBullet::Dash, "[] Not a checkbox"), + @" • [] Not a checkbox" + ); + } + + #[test] + fn checkbox_pattern_mid_content() { + insta::assert_snapshot!( + render(0, ListBullet::Dash, "Item with [ ] in middle"), + @" • Item with [ ] in middle" + ); + } + + #[test] + fn array_index_syntax() { + insta::assert_snapshot!( + render(0, ListBullet::Dash, "array[x] access"), + @" • array[x] access" + ); + } + + #[test] + fn map_access_syntax() { + insta::assert_snapshot!( + render(0, ListBullet::Dash, "map[ ] access"), + @" • map[ ] access" + ); + } + } + } +} diff --git a/crates/forge_markdown_stream/src/renderer.rs b/crates/forge_markdown_stream/src/renderer.rs new file mode 100644 index 0000000000000000000000000000000000000000..7b4962c979449b456c83daaa3f5e032d5b87991b --- /dev/null +++ b/crates/forge_markdown_stream/src/renderer.rs @@ -0,0 +1,315 @@ +//! Main renderer that handles all parse events. + +use std::io::{self, Write}; + +use streamdown_ansi::utils::visible_length; +use streamdown_parser::ParseEvent; + +use crate::code::CodeHighlighter; +use crate::heading::render_heading; +use crate::inline::{render_inline_content, render_inline_elements}; +use crate::list::{ListState, render_list_item}; +use crate::style::InlineStyler; +use crate::table::render_table; +use crate::theme::Theme; +use crate::utils::wrap_text_preserving_spaces; + +/// Main renderer for markdown events. +pub struct Renderer { + writer: W, + width: usize, + theme: Theme, + // Code highlighting + highlighter: CodeHighlighter, + current_language: Option, + code_buffer: String, + // Table buffering + table_rows: Vec>, + // Blockquote state + in_blockquote: bool, + blockquote_depth: usize, + // List state + list_state: ListState, + // Column tracking + column: usize, +} + +impl Renderer { + pub fn new(writer: W, width: usize) -> Self { + Self::with_theme(writer, width, Theme::default()) + } + + pub fn with_theme(writer: W, width: usize, theme: Theme) -> Self { + Self { + writer, + width, + theme, + highlighter: CodeHighlighter::default(), + current_language: None, + code_buffer: String::new(), + table_rows: Vec::new(), + in_blockquote: false, + blockquote_depth: 0, + list_state: ListState::default(), + column: 0, + } + } + + /// Set a new theme. + #[allow(dead_code)] + pub fn set_theme(&mut self, theme: Theme) { + self.theme = theme; + } + + /// Get the current theme. + #[allow(dead_code)] + pub fn theme(&self) -> &Theme { + &self.theme + } + + /// Calculate the left margin based on blockquote depth. + fn left_margin(&self) -> String { + if self.in_blockquote { + let border = self.theme.blockquote_border.apply("│").to_string(); + format!("{} ", border).repeat(self.blockquote_depth) + } else { + String::new() + } + } + + /// Calculate the current available width. + fn current_width(&self) -> usize { + let margin_width = if self.in_blockquote { + self.blockquote_depth * 3 + } else { + 0 + }; + self.width.saturating_sub(margin_width) + } + + fn write(&mut self, s: &str) -> io::Result<()> { + write!(self.writer, "{}", s) + } + + fn writeln(&mut self, s: &str) -> io::Result<()> { + writeln!(self.writer, "{}", s)?; + self.column = 0; + Ok(()) + } + + fn flush_table(&mut self) -> io::Result<()> { + if self.table_rows.is_empty() { + return Ok(()); + } + let rows = std::mem::take(&mut self.table_rows); + let margin = self.left_margin(); + let lines = render_table(&rows, &margin, &self.theme, self.width); + for line in lines { + self.writeln(&line)?; + } + Ok(()) + } + + /// Check if this event should reset a pending list. + /// List continues only for ListItem, ListEnd, and EmptyLine/Newline events. + fn should_reset_list(event: &ParseEvent) -> bool { + !matches!( + event, + ParseEvent::ListItem { .. } + | ParseEvent::ListEnd + | ParseEvent::EmptyLine + | ParseEvent::Newline + ) + } + + /// Render a single parse event. + pub fn render_event(&mut self, event: &ParseEvent) -> io::Result<()> { + // Reset pending list if this event breaks the list context + if Self::should_reset_list(event) { + self.list_state.reset(); + } + + match event { + // === Inline elements === + ParseEvent::Text(text) => { + let styled = self.theme.text(text); + self.write(&styled)?; + self.column += styled.chars().count(); + } + + ParseEvent::InlineCode(code) => { + self.write(&self.theme.code(code))?; + } + + ParseEvent::Bold(text) => { + self.write(&self.theme.bold(text))?; + } + + ParseEvent::Italic(text) => { + self.write(&self.theme.italic(text))?; + } + + ParseEvent::BoldItalic(text) => { + self.write(&self.theme.bold_italic(text))?; + } + + ParseEvent::Underline(text) => { + self.write(&self.theme.underline(text))?; + } + + ParseEvent::Strikeout(text) => { + self.write(&self.theme.strikethrough(text))?; + } + + ParseEvent::Link { text, url } => { + self.write(&self.theme.link(text, url))?; + } + + ParseEvent::Image { alt, url } => { + self.write(&self.theme.image(alt, url))?; + } + + ParseEvent::Footnote(superscript) => { + self.write(&self.theme.footnote(superscript))?; + } + + ParseEvent::Prompt(prompt) => { + self.write(prompt)?; + } + + // === Block elements === + ParseEvent::Heading { level, content } => { + let margin = self.left_margin(); + let width = self.current_width(); + let lines = render_heading(*level, content, width, &margin, &self.theme); + for line in lines { + self.writeln(&line)?; + } + } + + ParseEvent::CodeBlockStart { language, .. } => { + self.current_language = language.clone(); + self.code_buffer.clear(); + } + + ParseEvent::CodeBlockLine(line) => { + if !self.code_buffer.is_empty() { + self.code_buffer.push('\n'); + } + self.code_buffer.push_str(line); + + let margin = self.left_margin(); + let width = self.current_width(); + let rendered_lines = self.highlighter.render_code_line( + line, + self.current_language.as_deref(), + &margin, + width, + ); + for rendered in rendered_lines { + self.writeln(&rendered)?; + } + } + + ParseEvent::CodeBlockEnd => { + self.current_language = None; + self.code_buffer.clear(); + } + + ParseEvent::ListItem { indent, bullet, content } => { + let margin = self.left_margin(); + let width = self.current_width(); + let lines = render_list_item( + *indent, + bullet, + content, + width, + &margin, + &self.theme, + &mut self.list_state, + ); + for line in lines { + self.writeln(&line)?; + } + } + + ParseEvent::ListEnd => { + // Mark as pending - will reset if non-list event arrives + self.list_state.mark_pending_reset(); + } + + ParseEvent::TableHeader(cols) | ParseEvent::TableRow(cols) => { + self.table_rows.push(cols.clone()); + } + + ParseEvent::TableSeparator => {} + + ParseEvent::TableEnd => { + self.flush_table()?; + } + + ParseEvent::BlockquoteStart { depth } => { + self.in_blockquote = true; + self.blockquote_depth = *depth; + } + + ParseEvent::BlockquoteLine(text) => { + let margin = self.left_margin(); + let content_width = self.width.saturating_sub(visible_length(&margin)); + // Parse inline formatting (bold, italic, etc.) in blockquote content + let rendered_content = render_inline_content(text, &self.theme); + let wrapped = wrap_text_preserving_spaces( + &rendered_content, + content_width, + content_width, + &margin, + &margin, + ); + if wrapped.is_empty() { + self.writeln(&margin)?; + } else { + for line in wrapped { + self.writeln(&line)?; + } + } + } + + ParseEvent::BlockquoteEnd => { + self.in_blockquote = false; + self.blockquote_depth = 0; + } + + ParseEvent::ThinkBlockStart => { + self.writeln(&self.theme.think_border.apply("┌─ thinking ─").to_string())?; + self.in_blockquote = true; + self.blockquote_depth = 1; + } + + ParseEvent::ThinkBlockLine(text) => { + let border = self.theme.think_border.apply("│").to_string(); + self.writeln(&format!("{} {}", border, self.theme.think.apply(text)))?; + } + + ParseEvent::ThinkBlockEnd => { + self.writeln(&self.theme.think_border.apply("└").to_string())?; + self.in_blockquote = false; + self.blockquote_depth = 0; + } + + ParseEvent::HorizontalRule => { + let margin = self.left_margin(); + let rule = "─".repeat(self.current_width()); + self.writeln(&format!("{}{}", margin, self.theme.hr.apply(&rule)))?; + } + + ParseEvent::EmptyLine | ParseEvent::Newline => { + self.writeln("")?; + } + ParseEvent::InlineElements(elements) => { + self.write(&render_inline_elements(elements, &self.theme))?; + } + } + + self.writer.flush() + } +} diff --git a/crates/forge_markdown_stream/src/repair.rs b/crates/forge_markdown_stream/src/repair.rs new file mode 100644 index 0000000000000000000000000000000000000000..857d0fde19e561c8bf42e685223299cce0949b3a --- /dev/null +++ b/crates/forge_markdown_stream/src/repair.rs @@ -0,0 +1,123 @@ +//! Repair malformed markdown before parsing. +//! +//! This module handles common markdown issues that the parser doesn't handle +//! well, such as closing code fences on the same line as content. + +use streamdown_core::ParseState; + +/// Repair a line of markdown, returning one or more normalized lines. +/// +/// Handles: +/// - Embedded closing fences: `}```\n` becomes `}\n` + ```` ``` ```` (only when +/// in code block) +pub fn repair_line(line: &str, state: &ParseState) -> Vec { + // Only check for embedded closing fence when we're inside a code block + if state.is_in_code() + && let Some(lines) = split_embedded_fence(line) + { + return lines; + } + + vec![line.to_string()] +} + +/// Split a line if it contains an embedded closing fence at the end. +/// e.g., `}``` ` becomes Some(vec![`}`, ```` ``` ````]) +fn split_embedded_fence(line: &str) -> Option> { + let trimmed = line.trim_end(); + + // Check for ``` at the end + if let Some(stripped) = trimmed.strip_suffix("```") + && !stripped.trim().is_empty() + { + return Some(vec![stripped.to_string(), "```".to_string()]); + } + + // Check for ~~~ at the end + if let Some(stripped) = trimmed.strip_suffix("~~~") + && !stripped.trim().is_empty() + { + return Some(vec![stripped.to_string(), "~~~".to_string()]); + } + + None +} + +#[cfg(test)] +mod tests { + use streamdown_core::Code; + + use super::*; + + fn state_outside_code() -> ParseState { + ParseState::new() + } + + fn state_inside_code() -> ParseState { + let mut state = ParseState::new(); + state.enter_code_block(Code::Backtick, Some("rust".to_string())); + state + } + + #[test] + fn test_normal_line_unchanged() { + assert_eq!( + repair_line("hello world", &state_outside_code()), + vec!["hello world"] + ); + assert_eq!( + repair_line("hello world", &state_inside_code()), + vec!["hello world"] + ); + } + + #[test] + fn test_valid_fence_unchanged() { + assert_eq!(repair_line("```", &state_outside_code()), vec!["```"]); + assert_eq!(repair_line(" ```", &state_outside_code()), vec![" ```"]); + assert_eq!( + repair_line("```rust", &state_outside_code()), + vec!["```rust"] + ); + } + + #[test] + fn test_embedded_fence_not_split_outside_code_block() { + // Outside code block, don't split + assert_eq!(repair_line("}```", &state_outside_code()), vec!["}```"]); + assert_eq!( + repair_line("return x;```", &state_outside_code()), + vec!["return x;```"] + ); + } + + #[test] + fn test_embedded_backtick_fence_split_in_code_block() { + // Inside code block, split embedded fences + assert_eq!(repair_line("}```", &state_inside_code()), vec!["}", "```"]); + assert_eq!( + repair_line(" }```", &state_inside_code()), + vec![" }", "```"] + ); + assert_eq!( + repair_line("return x;```", &state_inside_code()), + vec!["return x;", "```"] + ); + } + + #[test] + fn test_embedded_tilde_fence_split_in_code_block() { + assert_eq!(repair_line("}~~~", &state_inside_code()), vec!["}", "~~~"]); + assert_eq!( + repair_line("return x;~~~", &state_inside_code()), + vec!["return x;", "~~~"] + ); + } + + #[test] + fn test_whitespace_only_before_fence_unchanged() { + // Just whitespace before fence is a valid fence, don't split + assert_eq!(repair_line(" ```", &state_inside_code()), vec![" ```"]); + assert_eq!(repair_line("\t```", &state_inside_code()), vec!["\t```"]); + } +} diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__all_empty_cells.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__all_empty_cells.snap new file mode 100644 index 0000000000000000000000000000000000000000..84dc2bff95c7c45f89746dc7472919187681d397 --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__all_empty_cells.snap @@ -0,0 +1,9 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render(vec![vec![\"\", \"\"], vec![\"\", \"\"],])" +--- + ┌──┬──┐ + │ │ │ + ├──┼──┤ + │ │ │ + └──┴──┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__cell_with_bold_tag.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__cell_with_bold_tag.snap new file mode 100644 index 0000000000000000000000000000000000000000..0d6b2d6526ba51585fb8c8782ae58b279a15000f --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__cell_with_bold_tag.snap @@ -0,0 +1,9 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render(vec![vec![\"Feature\", \"Status\"], vec![\"**Important**\", \"Done\"],])" +--- + ┌──────────────────┬────────┐ + │ Feature │ Status │ + ├──────────────────┼────────┤ + │ Important │ Done │ + └──────────────────┴────────┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__cell_with_code_tag.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__cell_with_code_tag.snap new file mode 100644 index 0000000000000000000000000000000000000000..22c23ce896124f8c53fa3ab12fd3fbad4711a9f5 --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__cell_with_code_tag.snap @@ -0,0 +1,11 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render(vec![vec![\"Function\", \"Description\"],\nvec![\"`render()`\", \"Renders the output\"],\nvec![\"`parse()`\", \"Parses input data\"],])" +--- + ┌───────────────────────┬────────────────────┐ + │ Function │ Description │ + ├───────────────────────┼────────────────────┤ + │ render() │ Renders the output │ + ├───────────────────────┼────────────────────┤ + │ parse() │ Parses input data │ + └───────────────────────┴────────────────────┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__cell_with_html_like_tags.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__cell_with_html_like_tags.snap new file mode 100644 index 0000000000000000000000000000000000000000..1a27ed22c54438d0b474f2e2087e1a8c0a82b6c2 --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__cell_with_html_like_tags.snap @@ -0,0 +1,11 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render(vec![vec![\"HTML\", \"Content\"], vec![\"
\", \"element\"],\nvec![\"\", \"styled\"],])" +--- + ┌──────────────────┬─────────┐ + │ HTML │ Content │ + ├──────────────────┼─────────┤ + │
│ element │ + ├──────────────────┼─────────┤ + │ │ styled │ + └──────────────────┴─────────┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__cell_with_italic_tag.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__cell_with_italic_tag.snap new file mode 100644 index 0000000000000000000000000000000000000000..b3b66141afc546fc0f4a2151fe26ba74bbf5fb79 --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__cell_with_italic_tag.snap @@ -0,0 +1,9 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render(vec![vec![\"Note\", \"Details\"], vec![\"*emphasis*\", \"Regular text\"],])" +--- + ┌─────────────────┬──────────────┐ + │ Note │ Details │ + ├─────────────────┼──────────────┤ + │ emphasis │ Regular text │ + └─────────────────┴──────────────┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__cell_with_link_tag.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__cell_with_link_tag.snap new file mode 100644 index 0000000000000000000000000000000000000000..bd27d8eecec793e70541608706c0ff195569d3c1 --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__cell_with_link_tag.snap @@ -0,0 +1,9 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render(vec![vec![\"Resource\", \"URL\"],\nvec![\"Documentation\", \"[docs](https://example.com)\"],])" +--- + ┌───────────────┬────────────────────────────────────────┐ + │ Resource │ URL │ + ├───────────────┼────────────────────────────────────────┤ + │ Documentation │ docs │ + └───────────────┴────────────────────────────────────────┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__cell_with_mixed_tags.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__cell_with_mixed_tags.snap new file mode 100644 index 0000000000000000000000000000000000000000..cd1bc15bd958d1df8b83eb24afb75ab46375f8eb --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__cell_with_mixed_tags.snap @@ -0,0 +1,11 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render(vec![vec![\"Item\", \"Description\"],\nvec![\"**Bold** and *italic*\", \"Mixed `code` here\"],\nvec![\"Normal\", \"[link](url) text\"],])" +--- + ┌───────────────────────────────┬──────────────────────────────┐ + │ Item │ Description │ + ├───────────────────────────────┼──────────────────────────────┤ + │ Bold and italic │ Mixed code here │ + ├───────────────────────────────┼──────────────────────────────┤ + │ Normal │ link text │ + └───────────────────────────────┴──────────────────────────────┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__cell_with_nested_tags.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__cell_with_nested_tags.snap new file mode 100644 index 0000000000000000000000000000000000000000..0fc1ae3488eb67a7d95e69dc38c3d888c6fbb78d --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__cell_with_nested_tags.snap @@ -0,0 +1,9 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render(vec![vec![\"Complex\", \"Content\"],\nvec![\"***bold italic***\", \"~~strikethrough~~\"],])" +--- + ┌───────────────────────────┬──────────────────────┐ + │ Complex │ Content │ + ├───────────────────────────┼──────────────────────┤ + │ bold italic │ strikethrough │ + └───────────────────────────┴──────────────────────┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__custom_margin.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__custom_margin.snap new file mode 100644 index 0000000000000000000000000000000000000000..70b0f797b7fe2332aec3282b7352ef2c9adb5d71 --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__custom_margin.snap @@ -0,0 +1,9 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render_with_margin(vec![vec![\"A\", \"B\"], vec![\"1\", \"2\"],], \" \")" +--- + ┌───┬───┐ + │ A │ B │ + ├───┼───┤ + │ 1 │ 2 │ + └───┴───┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__empty_cells.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__empty_cells.snap new file mode 100644 index 0000000000000000000000000000000000000000..7cbcfedd97d9159324446fa4daa7ac7b718c9fd1 --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__empty_cells.snap @@ -0,0 +1,11 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render(vec![vec![\"A\", \"B\"], vec![\"\", \"value\"], vec![\"data\", \"\"],])" +--- + ┌──────┬───────┐ + │ A │ B │ + ├──────┼───────┤ + │ │ value │ + ├──────┼───────┤ + │ data │ │ + └──────┴───────┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__header_with_tags.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__header_with_tags.snap new file mode 100644 index 0000000000000000000000000000000000000000..c035f55a78eb5027fd067cae77128fe6f9b967da --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__header_with_tags.snap @@ -0,0 +1,9 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render(vec![vec![\"**Bold Header**\", \"`Code Header`\"],\nvec![\"data1\", \"data2\"],])" +--- + ┌────────────────────┬──────────────────────────┐ + │ Bold Header │ Code Header │ + ├────────────────────┼──────────────────────────┤ + │ data1 │ data2 │ + └────────────────────┴──────────────────────────┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__long_content_in_cells.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__long_content_in_cells.snap new file mode 100644 index 0000000000000000000000000000000000000000..f5cfc9e4bc71671ba4ef903ca07db0f40027b2d9 --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__long_content_in_cells.snap @@ -0,0 +1,12 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +assertion_line: 564 +expression: "render(vec![vec![\"Header\", \"Description\"],\nvec![\"Short\",\n\"This is a much longer piece of content that should demonstrate how the table handles varying content lengths\"],])" +--- + ┌───────┬────────────────────────────────────────────────────────────────────┐ + │ Heade │ Description │ + │ r │ │ + ├───────┼────────────────────────────────────────────────────────────────────┤ + │ Short │ This is a much longer piece of content that should demonstrate │ + │ │ how the table handles varying content lengths │ + └───────┴────────────────────────────────────────────────────────────────────┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__long_content_with_tags_wrapping.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__long_content_with_tags_wrapping.snap new file mode 100644 index 0000000000000000000000000000000000000000..1b2bfc4a374aa4fa638783c00484ddb8ab919670 --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__long_content_with_tags_wrapping.snap @@ -0,0 +1,12 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +assertion_line: 731 +expression: "render_with_width(vec![vec![\"Title\", \"Content\"],\nvec![\"Article\",\n\"This has **bold** and *italic* and `code` in a long sentence that wraps\"],],\n50)" +--- + ┌───────┬──────────────────────────────────────┐ + │ Title │ Content │ + ├───────┼──────────────────────────────────────┤ + │ Artic │ This has bold and │ + │ le │ italic and code │ + │ │ in a long sentence that wraps │ + └───────┴──────────────────────────────────────┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__many_rows.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__many_rows.snap new file mode 100644 index 0000000000000000000000000000000000000000..ca8d371dd3b06e1232bd5f2da697f58b3e0bf940 --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__many_rows.snap @@ -0,0 +1,15 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render(vec![vec![\"ID\", \"Value\"], vec![\"1\", \"one\"], vec![\"2\", \"two\"],\nvec![\"3\", \"three\"], vec![\"4\", \"four\"],])" +--- + ┌────┬───────┐ + │ ID │ Value │ + ├────┼───────┤ + │ 1 │ one │ + ├────┼───────┤ + │ 2 │ two │ + ├────┼───────┤ + │ 3 │ three │ + ├────┼───────┤ + │ 4 │ four │ + └────┴───────┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__multiline_content_wrapping.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__multiline_content_wrapping.snap new file mode 100644 index 0000000000000000000000000000000000000000..44c3536ca254631e177982dc620300103b567e79 --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__multiline_content_wrapping.snap @@ -0,0 +1,14 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +assertion_line: 575 +expression: "render_with_width(vec![vec![\"Name\", \"Bio\"],\nvec![\"Alice\", \"Software engineer with 10 years of experience\"],\nvec![\"Bob\", \"Data scientist specializing in machine learning\"],], 50)" +--- + ┌───────┬──────────────────────────────────────┐ + │ Name │ Bio │ + ├───────┼──────────────────────────────────────┤ + │ Alice │ Software engineer with 10 years of │ + │ │ experience │ + ├───────┼──────────────────────────────────────┤ + │ Bob │ Data scientist specializing in │ + │ │ machine learning │ + └───────┴──────────────────────────────────────┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__narrow_first_col_not_padded.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__narrow_first_col_not_padded.snap new file mode 100644 index 0000000000000000000000000000000000000000..c77548228ad3db6b0ad246b6ed36adabc6c233ea --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__narrow_first_col_not_padded.snap @@ -0,0 +1,14 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +assertion_line: 499 +expression: "render_with_width(vec![vec![\"#\", \"File\", \"Description\"],\nvec![\"1\", \"plans/foo.md\",\n\"A longer description that forces the table to shrink its columns to fit the target width\"],\nvec![\"2\", \"docs/bar.md\",\n\"Another reasonably long description so the Description column stays the widest\"],],\n80)" +--- + ┌───┬──────────┬─────────────────────────────────────────────────────────────┐ + │ # │ File │ Description │ + ├───┼──────────┼─────────────────────────────────────────────────────────────┤ + │ 1 │ plans/fo │ A longer description that forces the table to shrink its │ + │ │ o.md │ columns to fit the target width │ + ├───┼──────────┼─────────────────────────────────────────────────────────────┤ + │ 2 │ docs/bar │ Another reasonably long description so the Description │ + │ │ .md │ column stays the widest │ + └───┴──────────┴─────────────────────────────────────────────────────────────┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__narrow_width_shrinks_columns.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__narrow_width_shrinks_columns.snap new file mode 100644 index 0000000000000000000000000000000000000000..14430fb7158fc06d6435c8e70dc8b2103bd2caba --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__narrow_width_shrinks_columns.snap @@ -0,0 +1,9 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render_with_width(vec![vec![\"Long Header One\", \"Long Header Two\"],\nvec![\"value1\", \"value2\"],], 40)" +--- + ┌─────────────────┬─────────────────┐ + │ Long Header One │ Long Header Two │ + ├─────────────────┼─────────────────┤ + │ value1 │ value2 │ + └─────────────────┴─────────────────┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__no_margin.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__no_margin.snap new file mode 100644 index 0000000000000000000000000000000000000000..4287d80f40db240bad5b8ed50676938d5b3b75e1 --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__no_margin.snap @@ -0,0 +1,9 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render_with_margin(vec![vec![\"A\", \"B\"], vec![\"1\", \"2\"],], \"\")" +--- +┌───┬───┐ +│ A │ B │ +├───┼───┤ +│ 1 │ 2 │ +└───┴───┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__numeric_content.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__numeric_content.snap new file mode 100644 index 0000000000000000000000000000000000000000..048319fd630e03ef56b892a059fc182975649d99 --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__numeric_content.snap @@ -0,0 +1,13 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render(vec![vec![\"ID\", \"Price\", \"Quantity\"], vec![\"1001\", \"$19.99\", \"150\"],\nvec![\"1002\", \"$249.50\", \"25\"], vec![\"1003\", \"$5.00\", \"1000\"],])" +--- + ┌──────┬─────────┬──────────┐ + │ ID │ Price │ Quantity │ + ├──────┼─────────┼──────────┤ + │ 1001 │ $19.99 │ 150 │ + ├──────┼─────────┼──────────┤ + │ 1002 │ $249.50 │ 25 │ + ├──────┼─────────┼──────────┤ + │ 1003 │ $5.00 │ 1000 │ + └──────┴─────────┴──────────┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__real_world_table_with_all_formatting.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__real_world_table_with_all_formatting.snap new file mode 100644 index 0000000000000000000000000000000000000000..81952e9788e70072b3513736b36b4eab03541ee1 --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__real_world_table_with_all_formatting.snap @@ -0,0 +1,32 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render(vec![vec![\"Feature\", \"Status\", \"Description\", \"Link\"],\nvec![\"**Authentication**\", \"✅ `completed`\",\n\"Implements *JWT-based* authentication with ~~basic~~ **OAuth2** support\",\n\"[Docs](https://example.com)\",],\nvec![\"**Database Layer**\", \"🚧 `in-progress`\",\n\"Uses `PostgreSQL` with **Diesel ORM** for *type-safe* queries\",\n\"[GitHub](https://github.com)\",],\nvec![\"**API Gateway**\", \"⏳ `planned`\",\n\"RESTful API with `async/await` and ~~synchronous~~ **asynchronous** handlers\",\n\"[Spec](https://api.example.com)\",],\nvec![\"**Testing**\", \"✅ `completed`\",\n\"Includes *unit tests*, **integration tests**, and `snapshot testing`\",\n\"[Coverage](https://coverage.io)\",],\nvec![\"**Deployment**\", \"🚧 `in-progress`\",\n\"Docker containerization with `K8s` orchestration and **CI/CD** pipeline\",\n\"[Deploy](https://deploy.com)\",],])" +--- + ┌─────────┬───────────┬───────────────────────────────────┬────────────────┐ + │ Feature │ Status │ Description │ Link │ + ├─────────┼───────────┼───────────────────────────────────┼────────────────┤ + │ Auth │ ✅ │ Implements JWT-based │ com │ authentication with basic │ href="https:// │ + │ ion │ pletedOAuth2 support │ example.com">D │ + │ │ ode> │ │ ocs │ + ├─────────┼───────────┼───────────────────────────────────┼────────────────┤ + │ Data │ 🚧 │ Uses PostgreSQL │ in- │ with Diesel ORM for │ href="https:// │ + │ Layertype-safe queries │ github.com">Gi │ + │ b> │ /code> │ │ tHub │ + ├─────────┼───────────┼───────────────────────────────────┼────────────────┤ + │ API │ ⏳ │ RESTful API with │ pla │ async/await and │ href="https:// │ + │ │ nnedsynchronous │ api.example.co │ + │ │ e> │ asynchronous handlers │ m">Spec │ + ├─────────┼───────────┼───────────────────────────────────┼────────────────┤ + │ Test │ ✅ │ Includes unit tests, │ │ com │ integration tests, and │ href="https:// │ + │ │ pletedsnapshot testing │ coverage.io">C │ + │ │ ode> │ │ overage │ + ├─────────┼───────────┼───────────────────────────────────┼────────────────┤ + │ Depl │ 🚧 │ Docker containerization with │ in- │ K8s orchestration │ href="https:// │ + │ /b> │ progress< │ and CI/CD pipeline │ deploy.com">De │ + │ │ /code> │ │ ploy │ + └─────────┴───────────┴───────────────────────────────────┴────────────────┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__simple_table.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__simple_table.snap new file mode 100644 index 0000000000000000000000000000000000000000..d8e72f72c427591ab1f586f3f57b42fac109817b --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__simple_table.snap @@ -0,0 +1,9 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render(vec![vec![\"Name\", \"Age\"], vec![\"Alice\", \"30\"],])" +--- + ┌───────┬─────┐ + │ Name │ Age │ + ├───────┼─────┤ + │ Alice │ 30 │ + └───────┴─────┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__single_cell.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__single_cell.snap new file mode 100644 index 0000000000000000000000000000000000000000..a25003cd80029eff350bafc9e0e7f286979930f9 --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__single_cell.snap @@ -0,0 +1,9 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render(vec![vec![\"Header\"], vec![\"Value\"],])" +--- + ┌────────┐ + │ Header │ + ├────────┤ + │ Value │ + └────────┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__single_column_many_rows.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__single_column_many_rows.snap new file mode 100644 index 0000000000000000000000000000000000000000..f062b5bfefb465b8be389527ea3c861317e27742 --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__single_column_many_rows.snap @@ -0,0 +1,17 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render(vec![vec![\"Items\"], vec![\"First\"], vec![\"Second\"], vec![\"Third\"],\nvec![\"Fourth\"], vec![\"Fifth\"],])" +--- + ┌────────┐ + │ Items │ + ├────────┤ + │ First │ + ├────────┤ + │ Second │ + ├────────┤ + │ Third │ + ├────────┤ + │ Fourth │ + ├────────┤ + │ Fifth │ + └────────┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__single_row_header_only.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__single_row_header_only.snap new file mode 100644 index 0000000000000000000000000000000000000000..cd773a69aa6597168657d46419fd8fba48932a2f --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__single_row_header_only.snap @@ -0,0 +1,7 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render(vec![vec![\"Only\", \"Headers\"],])" +--- + ┌──────┬─────────┐ + │ Only │ Headers │ + └──────┴─────────┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__special_characters.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__special_characters.snap new file mode 100644 index 0000000000000000000000000000000000000000..5aeb829e21c3c954b781a68d799935866f60e6fc --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__special_characters.snap @@ -0,0 +1,15 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render(vec![vec![\"Symbol\", \"Meaning\"], vec![\"<\", \"less than\"],\nvec![\">\", \"greater than\"], vec![\"&\", \"ampersand\"], vec![\"\\\"\", \"quote\"],])" +--- + ┌────────┬──────────────┐ + │ Symbol │ Meaning │ + ├────────┼──────────────┤ + │ < │ less than │ + ├────────┼──────────────┤ + │ > │ greater than │ + ├────────┼──────────────┤ + │ & │ ampersand │ + ├────────┼──────────────┤ + │ " │ quote │ + └────────┴──────────────┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__three_columns.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__three_columns.snap new file mode 100644 index 0000000000000000000000000000000000000000..a7da9de07b8989cf7fdf59f8d9851dc0562b2349 --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__three_columns.snap @@ -0,0 +1,11 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render(vec![vec![\"A\", \"B\", \"C\"], vec![\"1\", \"2\", \"3\"], vec![\"x\", \"y\", \"z\"],])" +--- + ┌───┬───┬───┐ + │ A │ B │ C │ + ├───┼───┼───┤ + │ 1 │ 2 │ 3 │ + ├───┼───┼───┤ + │ x │ y │ z │ + └───┴───┴───┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__uneven_row_lengths.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__uneven_row_lengths.snap new file mode 100644 index 0000000000000000000000000000000000000000..fbbab79aff56a39b219b1de37e96ec1dd49423ca --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__uneven_row_lengths.snap @@ -0,0 +1,11 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render(vec![vec![\"A\", \"B\", \"C\"], vec![\"1\", \"2\"], vec![\"x\"],])" +--- + ┌───┬───┬───┐ + │ A │ B │ C │ + ├───┼───┼───┤ + │ 1 │ 2 │ │ + ├───┼───┼───┤ + │ x │ │ │ + └───┴───┴───┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__unicode_content.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__unicode_content.snap new file mode 100644 index 0000000000000000000000000000000000000000..818822539188942a7f787fe8fb14b1cb8b5edf68 --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__unicode_content.snap @@ -0,0 +1,9 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render(vec![vec![\"名前\", \"年齢\"], vec![\"田中\", \"25\"],])" +--- + ┌──────┬──────┐ + │ 名前 │ 年齢 │ + ├──────┼──────┤ + │ 田中 │ 25 │ + └──────┴──────┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__unicode_with_tags.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__unicode_with_tags.snap new file mode 100644 index 0000000000000000000000000000000000000000..f8e6d975fa647647b4403b6fa81afa15a81eff4d --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__unicode_with_tags.snap @@ -0,0 +1,11 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render(vec![vec![\"言語\", \"説明\"], vec![\"**日本語**\", \"*Japanese*\"],\nvec![\"`中文`\", \"Chinese\"],])" +--- + ┌───────────────────┬─────────────────┐ + │ 言語 │ 説明 │ + ├───────────────────┼─────────────────┤ + │ 日本語 │ Japanese │ + ├───────────────────┼─────────────────┤ + │ 中文 │ Chinese │ + └───────────────────┴─────────────────┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__varying_column_widths.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__varying_column_widths.snap new file mode 100644 index 0000000000000000000000000000000000000000..0d854f0c85a3358e761bc968e5c2f5b5e8089fd0 --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__varying_column_widths.snap @@ -0,0 +1,9 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render(vec![vec![\"Short\", \"Much Longer Header\"], vec![\"a\", \"b\"],])" +--- + ┌───────┬────────────────────┐ + │ Short │ Much Longer Header │ + ├───────┼────────────────────┤ + │ a │ b │ + └───────┴────────────────────┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__very_narrow_width.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__very_narrow_width.snap new file mode 100644 index 0000000000000000000000000000000000000000..e541044e61f1e18f5a6c0badfb18b950f8a54e22 --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__very_narrow_width.snap @@ -0,0 +1,9 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render_with_width(vec![vec![\"Column A\", \"Column B\"],\nvec![\"value1\", \"value2\"],], 25)" +--- + ┌──────────┬──────────┐ + │ Column A │ Column B │ + ├──────────┼──────────┤ + │ value1 │ value2 │ + └──────────┴──────────┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__whitespace_content.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__whitespace_content.snap new file mode 100644 index 0000000000000000000000000000000000000000..005ff4fd0585adf5ef2b0a69c9b0e4c681dcdd09 --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__whitespace_content.snap @@ -0,0 +1,9 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render(vec![vec![\"Header\", \"Value\"], vec![\" spaces \", \" \"],])" +--- + ┌────────────┬───────┐ + │ Header │ Value │ + ├────────────┼───────┤ + │ spaces │ │ + └────────────┴───────┘ diff --git a/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__wide_table_many_columns.snap b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__wide_table_many_columns.snap new file mode 100644 index 0000000000000000000000000000000000000000..1a10c58632660fb95b6ee2f8964184abd1b79ce0 --- /dev/null +++ b/crates/forge_markdown_stream/src/snapshots/forge_markdown_stream__table__tests__wide_table_many_columns.snap @@ -0,0 +1,9 @@ +--- +source: crates/forge_markdown_stream/src/table.rs +expression: "render(vec![vec![\"A\", \"B\", \"C\", \"D\", \"E\", \"F\"],\nvec![\"1\", \"2\", \"3\", \"4\", \"5\", \"6\"],])" +--- + ┌───┬───┬───┬───┬───┬───┐ + │ A │ B │ C │ D │ E │ F │ + ├───┼───┼───┼───┼───┼───┤ + │ 1 │ 2 │ 3 │ 4 │ 5 │ 6 │ + └───┴───┴───┴───┴───┴───┘ diff --git a/crates/forge_markdown_stream/src/style.rs b/crates/forge_markdown_stream/src/style.rs new file mode 100644 index 0000000000000000000000000000000000000000..666d93223da005aa7e10575b5ae1cbd7aaba420a --- /dev/null +++ b/crates/forge_markdown_stream/src/style.rs @@ -0,0 +1,41 @@ +/// Trait for styling inline elements. +pub trait InlineStyler { + fn text(&self, text: &str) -> String; + fn bold(&self, text: &str) -> String; + fn italic(&self, text: &str) -> String; + fn bold_italic(&self, text: &str) -> String; + fn strikethrough(&self, text: &str) -> String; + fn underline(&self, text: &str) -> String; + fn code(&self, text: &str) -> String; + fn link(&self, text: &str, url: &str) -> String; + fn image(&self, alt: &str, url: &str) -> String; + fn footnote(&self, text: &str) -> String; + fn dimmed(&self, text: &str) -> String; +} + +/// Trait for styling heading elements. +pub trait HeadingStyler { + fn h1(&self, text: &str) -> String; + fn h2(&self, text: &str) -> String; + fn h3(&self, text: &str) -> String; + fn h4(&self, text: &str) -> String; + fn h5(&self, text: &str) -> String; + fn h6(&self, text: &str) -> String; +} + +/// Trait for styling list elements. +pub trait ListStyler { + fn bullet_dash(&self, text: &str) -> String; + fn bullet_asterisk(&self, text: &str) -> String; + fn bullet_plus(&self, text: &str) -> String; + fn bullet_plus_expand(&self, text: &str) -> String; + fn number(&self, text: &str) -> String; + fn checkbox_checked(&self, text: &str) -> String; + fn checkbox_unchecked(&self, text: &str) -> String; +} + +/// Trait for styling table elements. +pub trait TableStyler { + fn border(&self, text: &str) -> String; + fn header(&self, text: &str) -> String; +} diff --git a/crates/forge_markdown_stream/src/table.rs b/crates/forge_markdown_stream/src/table.rs new file mode 100644 index 0000000000000000000000000000000000000000..0806974fd23496b03fa5127982688305812faebb --- /dev/null +++ b/crates/forge_markdown_stream/src/table.rs @@ -0,0 +1,833 @@ +//! Table rendering with box-drawing characters. + +use streamdown_ansi::utils::visible_length; +use unicode_width::{UnicodeWidthChar, UnicodeWidthStr}; + +use crate::inline::render_inline_content; +use crate::style::{InlineStyler, TableStyler}; + +/// Render a table with proper column widths, shrinking and wrapping if needed. +pub fn render_table( + rows: &[Vec], + margin: &str, + styler: &S, + max_width: usize, +) -> Vec { + // First, render all cells with inline markdown + let rendered_rows: Vec> = rows + .iter() + .map(|row| { + row.iter() + .map(|cell| render_inline_content(cell, styler)) + .collect() + }) + .collect(); + let n = rendered_rows.iter().map(|r| r.len()).max().unwrap_or(0); + if rendered_rows.is_empty() || n == 0 { + return vec![]; + } + + // Calculate column widths based on rendered content + let mut w: Vec = vec![0; n]; + for row in &rendered_rows { + for (i, cell) in row.iter().enumerate() { + let new_width = visible_length(cell); + if let Some(wi) = w.get_mut(i) { + *wi = (*wi).max(new_width); + } + } + } + + // Shrink columns if table exceeds max width + const MIN_COL_WIDTH: usize = 5; + let overhead = margin.width() + 1 + 3 * n; + let total: usize = w.iter().sum(); + if overhead + total > max_width && max_width > overhead { + let avail = max_width - overhead; + // Cap at natural width so narrow columns (e.g. "#") aren't inflated + // up to MIN_COL_WIDTH when the proportional share rounds below it. + w.iter_mut() + .for_each(|x| *x = (*x * avail / total).max(MIN_COL_WIDTH).min(*x)); + + // Clamping to MIN_COL_WIDTH can push the new total over `avail` when + // tiny columns get bumped up. Trim 1 char at a time from the widest + // column (above the minimum) until it fits. + let mut excess = w.iter().sum::().saturating_sub(avail); + while excess > 0 { + let Some(v) = w + .iter_mut() + .filter(|v| **v > MIN_COL_WIDTH) + .max_by_key(|v| **v) + else { + break; + }; + *v -= 1; + excess -= 1; + } + } + + // Helper to create horizontal lines + let hline = |l: &str, m: &str, r: &str| { + format!( + "{}{}{}{}", + margin, + styler.border(l), + w.iter() + .map(|&x| styler.border(&"─".repeat(x + 2))) + .collect::>() + .join(&styler.border(m)), + styler.border(r) + ) + }; + + let mut out = vec![hline("┌", "┬", "┐")]; + + for (ri, row) in rendered_rows.iter().enumerate() { + // Wrap each cell's content + let wrapped: Vec> = (0..n) + .map(|i| { + let width = w.get(i).copied().unwrap_or(0); + wrap(row.get(i).map(|s| s.as_str()).unwrap_or(""), width) + }) + .collect(); + + // Render each line of the wrapped cells + for li in 0..wrapped.iter().map(|c| c.len()).max().unwrap_or(1) { + let cells: String = (0..n) + .map(|i| { + let c = wrapped + .get(i) + .and_then(|w| w.get(li)) + .map(|s| s.as_str()) + .unwrap_or(""); + let width = w.get(i).copied().unwrap_or(0); + let p = " ".repeat(width.saturating_sub(visible_length(c))); + if ri == 0 && li == 0 && !c.is_empty() { + format!(" {}{} ", styler.header(c), p) + } else { + format!(" {}{} ", c, p) + } + }) + .collect::>() + .join(&styler.border("│")); + out.push(format!( + "{}{}{}{}", + margin, + styler.border("│"), + cells, + styler.border("│") + )); + } + + // Add row separator (except after last row) + if ri < rendered_rows.len() - 1 { + out.push(hline("├", "┼", "┤")); + } + } + + out.push(hline("└", "┴", "┘")); + out +} + +/// Wrap text by words, preserving ANSI codes across lines. +/// Breaks at spaces, and tries to keep content together when possible. +/// Handles both CSI sequences (\x1b[...m) and OSC sequences (\x1b]...\x1b\\). +fn wrap(text: &str, width: usize) -> Vec { + if width == 0 || visible_length(text) <= width { + return vec![text.to_string()]; + } + + let mut lines = Vec::new(); + let mut line = String::new(); + let mut line_width = 0; + let mut word = String::new(); + let mut word_width = 0; + let mut esc = String::new(); + let mut in_osc = false; + let mut active_style: Option = None; + + let chars: Vec = text.chars().collect(); + let mut i = 0; + + while i < chars.len() { + let c = chars.get(i).copied().unwrap_or('\0'); + + // Handle escape sequences + if c == '\x1b' { + esc.push(c); + i += 1; + + // Check what type of sequence + if i < chars.len() { + let next = chars.get(i).copied().unwrap_or('\0'); + esc.push(next); + i += 1; + + if next == '[' { + // CSI sequence - read until 'm' or other terminator + while i < chars.len() { + let sc = chars.get(i).copied().unwrap_or('\0'); + esc.push(sc); + i += 1; + if sc == 'm' || sc == 'K' || sc == 'H' || sc == 'J' { + break; + } + } + // Track active style for CSI color/style sequences + if esc.ends_with('m') { + active_style = if esc == "\x1b[0m" { + None + } else { + Some(esc.clone()) + }; + } + word.push_str(&esc); + esc.clear(); + } else if next == ']' { + // OSC sequence - read until \x1b\\ + in_osc = true; + while i < chars.len() { + let sc = chars.get(i).copied().unwrap_or('\0'); + esc.push(sc); + i += 1; + if sc == '\\' && esc.len() >= 2 { + let prev = esc.chars().rev().nth(1); + if prev == Some('\x1b') { + in_osc = false; + break; + } + } + // Also check for BEL terminator + if sc == '\x07' { + in_osc = false; + break; + } + } + word.push_str(&esc); + esc.clear(); + } else if next == '\\' && in_osc { + // End of OSC sequence + word.push_str(&esc); + esc.clear(); + in_osc = false; + } else { + // Unknown sequence, just add it + word.push_str(&esc); + esc.clear(); + } + } + continue; + } + + let cw = c.width().unwrap_or(0); + + // Check if this is a word boundary (space) + if c.is_whitespace() { + // Gate on `line_width` (visible), not `line.is_empty()`: when an + // active style is set `line` carries an ANSI prefix with zero + // visible chars, and pushing it would render as a blank row. + if line_width > 0 && line_width + word_width + cw > width { + line.push_str("\x1b[0m"); + lines.push(line); + line = active_style.clone().unwrap_or_default(); + line_width = 0; + } + line.push_str(&word); + line_width += word_width; + // Skip the separator if it would push past `width` — otherwise a + // word that fills the column exactly leaves a trailing space that + // overflows the cell by one char. + if line_width + cw <= width { + line.push(c); + line_width += cw; + } + word.clear(); + word_width = 0; + } else { + // Add character to current word + word.push(c); + word_width += cw; + + // If word itself exceeds width, break it by character + if word_width > width { + if line_width > 0 { + line.push_str("\x1b[0m"); + lines.push(line); + line = active_style.clone().unwrap_or_default(); + line_width = 0; + } + // Push the long word, breaking at width while preserving ANSI codes + while visible_length(&word) > width { + let (chunk, rem) = split_word_at_width(&word, width); + if !chunk.is_empty() { + line.push_str(&chunk); + line.push_str("\x1b[0m"); + lines.push(line); + line = active_style.clone().unwrap_or_default(); + line_width = 0; + } + word = rem; + } + word_width = visible_length(&word); + } + } + i += 1; + } + + // Add remaining word to line + if !word.is_empty() { + if line_width + word_width > width && line_width > 0 { + line.push_str("\x1b[0m"); + lines.push(line); + line = active_style.clone().unwrap_or_default(); + line.push_str(&word); + line_width = word_width; + } else { + line.push_str(&word); + line_width += word_width; + } + } + + if line_width > 0 { + lines.push(line); + } + + if lines.is_empty() { + vec![String::new()] + } else { + lines + } +} + +/// Split a word at a given visible width, preserving ANSI escape sequences. +fn split_word_at_width(word: &str, width: usize) -> (String, String) { + let mut chunk = String::new(); + let mut chunk_w = 0; + let mut remaining = String::new(); + let mut in_chunk = true; + let mut esc = String::new(); + + let chars: Vec = word.chars().collect(); + let mut i = 0; + + while i < chars.len() { + let c = chars.get(i).copied().unwrap_or('\0'); + + // Handle escape sequences + if c == '\x1b' { + esc.push(c); + i += 1; + + if i < chars.len() { + let next = chars.get(i).copied().unwrap_or('\0'); + esc.push(next); + i += 1; + + if next == '[' { + // CSI sequence + while i < chars.len() { + let sc = chars.get(i).copied().unwrap_or('\0'); + esc.push(sc); + i += 1; + if sc == 'm' || sc == 'K' || sc == 'H' || sc == 'J' { + break; + } + } + } else if next == ']' { + // OSC sequence - read until \x1b\\ or BEL + while i < chars.len() { + let sc = chars.get(i).copied().unwrap_or('\0'); + esc.push(sc); + i += 1; + if sc == '\\' && esc.len() >= 2 { + let prev_idx = esc.len() - 2; + if esc.chars().nth(prev_idx) == Some('\x1b') { + break; + } + } + if sc == '\x07' { + break; + } + } + } + } + + // Add escape sequence to appropriate part + if in_chunk { + chunk.push_str(&esc); + } else { + remaining.push_str(&esc); + } + esc.clear(); + continue; + } + + if in_chunk { + let cw = c.width().unwrap_or(0); + if chunk_w + cw <= width { + chunk.push(c); + chunk_w += cw; + } else { + remaining.push(c); + in_chunk = false; + } + } else { + remaining.push(c); + } + i += 1; + } + + (chunk, remaining) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::theme::{TagStyler, Theme}; + + fn strip_ansi(s: &str) -> String { + let bytes = strip_ansi_escapes::strip(s); + String::from_utf8(bytes).unwrap() + } + + fn render(rows: Vec>) -> String { + let rows: Vec> = rows + .into_iter() + .map(|r| r.into_iter().map(|s| s.to_string()).collect()) + .collect(); + let output = render_table(&rows, " ", &TagStyler, 80).join("\n"); + strip_ansi(&output) + } + + fn render_with_width(rows: Vec>, width: usize) -> String { + let rows: Vec> = rows + .into_iter() + .map(|r| r.into_iter().map(|s| s.to_string()).collect()) + .collect(); + let output = render_table(&rows, " ", &TagStyler, width).join("\n"); + strip_ansi(&output) + } + + fn render_with_margin(rows: Vec>, margin: &str) -> String { + let rows: Vec> = rows + .into_iter() + .map(|r| r.into_iter().map(|s| s.to_string()).collect()) + .collect(); + let output = render_table(&rows, margin, &TagStyler, 80).join("\n"); + strip_ansi(&output) + } + + #[test] + fn test_simple_table() { + insta::assert_snapshot!(render(vec![vec!["Name", "Age"], vec!["Alice", "30"],])); + } + + #[test] + fn test_single_cell() { + insta::assert_snapshot!(render(vec![vec!["Header"], vec!["Value"],])); + } + + #[test] + fn test_three_columns() { + insta::assert_snapshot!(render(vec![ + vec!["A", "B", "C"], + vec!["1", "2", "3"], + vec!["x", "y", "z"], + ])); + } + + #[test] + fn test_varying_column_widths() { + insta::assert_snapshot!(render(vec![ + vec!["Short", "Much Longer Header"], + vec!["a", "b"], + ])); + } + + #[test] + fn test_empty_cells() { + insta::assert_snapshot!(render(vec![ + vec!["A", "B"], + vec!["", "value"], + vec!["data", ""], + ])); + } + + #[test] + fn test_empty_table() { + let rows: Vec> = vec![]; + let result = render_table(&rows, " ", &Theme::dark(), 80); + assert!(result.is_empty()); + } + + #[test] + fn test_empty_row() { + let rows: Vec> = vec![vec![]]; + let result = render_table(&rows, " ", &Theme::dark(), 80); + assert!(result.is_empty()); + } + + #[test] + fn test_custom_margin() { + insta::assert_snapshot!(render_with_margin( + vec![vec!["A", "B"], vec!["1", "2"],], + " " + )); + } + + #[test] + fn test_no_margin() { + insta::assert_snapshot!(render_with_margin( + vec![vec!["A", "B"], vec!["1", "2"],], + "" + )); + } + + #[test] + fn test_narrow_width_shrinks_columns() { + insta::assert_snapshot!(render_with_width( + vec![ + vec!["Long Header One", "Long Header Two"], + vec!["value1", "value2"], + ], + 40 + )); + } + + #[test] + fn test_narrow_first_col_not_padded() { + insta::assert_snapshot!(render_with_width( + vec![ + vec!["#", "File", "Description"], + vec![ + "1", + "plans/foo.md", + "A longer description that forces the table to shrink its columns to fit the target width" + ], + vec![ + "2", + "docs/bar.md", + "Another reasonably long description so the Description column stays the widest" + ], + ], + 80 + )); + } + + #[test] + fn test_wrap_word_exactly_fills_width_drops_trailing_space() { + // A word whose width equals `width` followed by a space must not keep + // the space — it would push visible width to `width + 1` and overflow + // the cell border by one char. + let result = wrap("abcdefgh more text", 8); + let strip = |s: &str| String::from_utf8(strip_ansi_escapes::strip(s)).unwrap(); + assert_eq!(strip(&result[0]), "abcdefgh"); + for line in &result { + assert!(visible_length(line) <= 8); + } + } + + #[test] + fn test_styled_content_no_blank_middle_line() { + // An inherited ANSI style prefix on a fresh line has zero visible + // chars, so it must not be pushed as its own wrapped row. + let wrapped = wrap("\x1b[33mdocs/pdb-reference.md\x1b[0m", 8); + let stripped: Vec = wrapped + .iter() + .map(|s| String::from_utf8(strip_ansi_escapes::strip(s)).unwrap()) + .collect(); + assert_eq!(stripped, vec!["docs/pdb", "-referen", "ce.md"]); + } + + #[test] + fn test_unicode_content() { + insta::assert_snapshot!(render(vec![vec!["名前", "年齢"], vec!["田中", "25"],])); + } + + #[test] + fn test_single_row_header_only() { + insta::assert_snapshot!(render(vec![vec!["Only", "Headers"],])); + } + + #[test] + fn test_many_rows() { + insta::assert_snapshot!(render(vec![ + vec!["ID", "Value"], + vec!["1", "one"], + vec!["2", "two"], + vec!["3", "three"], + vec!["4", "four"], + ])); + } + + // ==================== wrap function tests ==================== + + #[test] + fn test_wrap_no_wrap_needed() { + let result = wrap("hello", 10); + assert_eq!(result, vec!["hello"]); + } + + #[test] + fn test_wrap_exact_width() { + let result = wrap("hello", 5); + assert_eq!(result, vec!["hello"]); + } + + #[test] + fn test_wrap_splits_text() { + let result = wrap("hello world", 5); + assert_eq!(result.len(), 2); + let strip = |s: &str| String::from_utf8(strip_ansi_escapes::strip(s)).unwrap(); + assert_eq!(strip(&result[0]), "hello"); + assert_eq!(strip(&result[1]), "world"); + } + + #[test] + fn test_wrap_empty() { + let result = wrap("", 10); + assert_eq!(result, vec![""]); + } + + #[test] + fn test_wrap_zero_width() { + let result = wrap("hello", 0); + assert_eq!(result, vec!["hello"]); + } + + #[test] + fn test_wrap_unicode() { + // Chinese chars are 2 wide each + let result = wrap("你好世界", 4); + assert_eq!(result.len(), 2); // "你好" and "世界" + } + + // ==================== Edge cases ==================== + + #[test] + fn test_long_content_in_cells() { + insta::assert_snapshot!(render(vec![ + vec!["Header", "Description"], + vec![ + "Short", + "This is a much longer piece of content that should demonstrate how the table handles varying content lengths" + ], + ])); + } + + #[test] + fn test_multiline_content_wrapping() { + insta::assert_snapshot!(render_with_width( + vec![ + vec!["Name", "Bio"], + vec!["Alice", "Software engineer with 10 years of experience"], + vec!["Bob", "Data scientist specializing in machine learning"], + ], + 50 + )); + } + + #[test] + fn test_very_narrow_width() { + insta::assert_snapshot!(render_with_width( + vec![vec!["Column A", "Column B"], vec!["value1", "value2"],], + 25 + )); + } + + #[test] + fn test_single_column_many_rows() { + insta::assert_snapshot!(render(vec![ + vec!["Items"], + vec!["First"], + vec!["Second"], + vec!["Third"], + vec!["Fourth"], + vec!["Fifth"], + ])); + } + + #[test] + fn test_wide_table_many_columns() { + insta::assert_snapshot!(render(vec![ + vec!["A", "B", "C", "D", "E", "F"], + vec!["1", "2", "3", "4", "5", "6"], + ])); + } + + #[test] + fn test_uneven_row_lengths() { + // Rows with different number of cells + insta::assert_snapshot!(render(vec![ + vec!["A", "B", "C"], + vec!["1", "2"], // Missing third cell + vec!["x"], // Only one cell + ])); + } + + #[test] + fn test_all_empty_cells() { + insta::assert_snapshot!(render(vec![vec!["", ""], vec!["", ""],])); + } + + #[test] + fn test_whitespace_content() { + insta::assert_snapshot!(render(vec![ + vec!["Header", "Value"], + vec![" spaces ", " "], + ])); + } + + #[test] + fn test_special_characters() { + insta::assert_snapshot!(render(vec![ + vec!["Symbol", "Meaning"], + vec!["<", "less than"], + vec![">", "greater than"], + vec!["&", "ampersand"], + vec!["\"", "quote"], + ])); + } + + #[test] + fn test_numeric_content() { + insta::assert_snapshot!(render(vec![ + vec!["ID", "Price", "Quantity"], + vec!["1001", "$19.99", "150"], + vec!["1002", "$249.50", "25"], + vec!["1003", "$5.00", "1000"], + ])); + } + + // ==================== Inner tags / inline formatting ==================== + + #[test] + fn test_cell_with_bold_tag() { + // Bold text in cell - tags should pass through as-is (not rendered) + insta::assert_snapshot!(render(vec![ + vec!["Feature", "Status"], + vec!["**Important**", "Done"], + ])); + } + + #[test] + fn test_cell_with_italic_tag() { + insta::assert_snapshot!(render(vec![ + vec!["Note", "Details"], + vec!["*emphasis*", "Regular text"], + ])); + } + + #[test] + fn test_cell_with_code_tag() { + insta::assert_snapshot!(render(vec![ + vec!["Function", "Description"], + vec!["`render()`", "Renders the output"], + vec!["`parse()`", "Parses input data"], + ])); + } + + #[test] + fn test_cell_with_link_tag() { + insta::assert_snapshot!(render(vec![ + vec!["Resource", "URL"], + vec!["Documentation", "[docs](https://example.com)"], + ])); + } + + #[test] + fn test_cell_with_mixed_tags() { + insta::assert_snapshot!(render(vec![ + vec!["Item", "Description"], + vec!["**Bold** and *italic*", "Mixed `code` here"], + vec!["Normal", "[link](url) text"], + ])); + } + + #[test] + fn test_cell_with_nested_tags() { + insta::assert_snapshot!(render(vec![ + vec!["Complex", "Content"], + vec!["***bold italic***", "~~strikethrough~~"], + ])); + } + + #[test] + fn test_header_with_tags() { + // Tags in header row + insta::assert_snapshot!(render(vec![ + vec!["**Bold Header**", "`Code Header`"], + vec!["data1", "data2"], + ])); + } + + #[test] + fn test_cell_with_html_like_tags() { + // HTML-like content should be preserved + insta::assert_snapshot!(render(vec![ + vec!["HTML", "Content"], + vec!["
", "element"], + vec!["", "styled"], + ])); + } + + #[test] + fn test_long_content_with_tags_wrapping() { + insta::assert_snapshot!(render_with_width( + vec![ + vec!["Title", "Content"], + vec![ + "Article", + "This has **bold** and *italic* and `code` in a long sentence that wraps" + ], + ], + 50 + )); + } + + #[test] + fn test_unicode_with_tags() { + insta::assert_snapshot!(render(vec![ + vec!["言語", "説明"], + vec!["**日本語**", "*Japanese*"], + vec!["`中文`", "Chinese"], + ])); + } + + #[test] + fn test_real_world_table_with_all_formatting() { + insta::assert_snapshot!(render(vec![ + vec!["Feature", "Status", "Description", "Link"], + vec![ + "**Authentication**", + "✅ `completed`", + "Implements *JWT-based* authentication with ~~basic~~ **OAuth2** support", + "[Docs](https://example.com)", + ], + vec![ + "**Database Layer**", + "🚧 `in-progress`", + "Uses `PostgreSQL` with **Diesel ORM** for *type-safe* queries", + "[GitHub](https://github.com)", + ], + vec![ + "**API Gateway**", + "⏳ `planned`", + "RESTful API with `async/await` and ~~synchronous~~ **asynchronous** handlers", + "[Spec](https://api.example.com)", + ], + vec![ + "**Testing**", + "✅ `completed`", + "Includes *unit tests*, **integration tests**, and `snapshot testing`", + "[Coverage](https://coverage.io)", + ], + vec![ + "**Deployment**", + "🚧 `in-progress`", + "Docker containerization with `K8s` orchestration and **CI/CD** pipeline", + "[Deploy](https://deploy.com)", + ], + ])); + } +} diff --git a/crates/forge_markdown_stream/src/theme.rs b/crates/forge_markdown_stream/src/theme.rs new file mode 100644 index 0000000000000000000000000000000000000000..ef2624a7a6d71cd31ba6c2318cfa2048c098076e --- /dev/null +++ b/crates/forge_markdown_stream/src/theme.rs @@ -0,0 +1,495 @@ +//! Theme configuration for markdown rendering. +//! +//! Provides customizable styling for all markdown elements using the `colored` +//! crate. + +use colored::{Color, ColoredString, Colorize}; +use streamdown_parser::decode_html_entities; + +use crate::style::{HeadingStyler, InlineStyler, ListStyler, TableStyler}; + +/// Style configuration for a single element. +#[derive(Clone, Debug, Default)] +pub struct Style { + pub fg: Option, + pub bg: Option, + pub bold: bool, + pub italic: bool, + pub underline: bool, + pub strikethrough: bool, + pub dimmed: bool, +} + +impl Style { + pub fn new() -> Self { + Self::default() + } + + pub fn fg(mut self, color: Color) -> Self { + self.fg = Some(color); + self + } + + pub fn bg(mut self, color: Color) -> Self { + self.bg = Some(color); + self + } + + pub fn bold(mut self) -> Self { + self.bold = true; + self + } + + pub fn italic(mut self) -> Self { + self.italic = true; + self + } + + pub fn underline(mut self) -> Self { + self.underline = true; + self + } + + pub fn strikethrough(mut self) -> Self { + self.strikethrough = true; + self + } + + pub fn dimmed(mut self) -> Self { + self.dimmed = true; + self + } + + /// Apply this style to a string. + pub fn apply(&self, text: &str) -> ColoredString { + let mut result = text.normal(); + + if let Some(fg) = self.fg { + result = result.color(fg); + } + if let Some(bg) = self.bg { + result = result.on_color(bg); + } + if self.bold { + result = result.bold(); + } + if self.italic { + result = result.italic(); + } + if self.underline { + result = result.underline(); + } + if self.strikethrough { + result = result.strikethrough(); + } + if self.dimmed { + result = result.dimmed(); + } + + result + } +} + +/// Theme containing styles for all markdown elements. +#[derive(Clone, Debug)] +pub struct Theme { + // Inline styles + pub bold: Style, + pub italic: Style, + pub code: Style, + pub strikethrough: Style, + pub link: Style, + pub link_url: Style, + + // Block styles + pub heading1: Style, + pub heading2: Style, + pub heading3: Style, + pub heading4: Style, + pub heading5: Style, + pub heading6: Style, + + // List styles + pub bullet_dash: Style, + pub bullet_asterisk: Style, + pub bullet_plus: Style, + pub bullet_plus_expand: Style, + pub list_number: Style, + pub checkbox_checked: Style, + pub checkbox_unchecked: Style, + + // Table styles + pub table_header: Style, + pub table_border: Style, + pub table_cell: Style, + + // Quote/Think styles + pub blockquote: Style, + pub blockquote_border: Style, + pub think: Style, + pub think_border: Style, + + // Code block + pub code_block_lang: Style, + + // Horizontal rule + pub hr: Style, +} + +impl Default for Theme { + fn default() -> Self { + Self::detect() + } +} + +/// Theme-based styler that outputs ANSI codes. +impl InlineStyler for Theme { + fn text(&self, text: &str) -> String { + decode_html_entities(text) + } + + fn bold(&self, text: &str) -> String { + self.bold.apply(&decode_html_entities(text)).to_string() + } + + fn italic(&self, text: &str) -> String { + self.italic.apply(&decode_html_entities(text)).to_string() + } + + fn bold_italic(&self, text: &str) -> String { + let decoded = decode_html_entities(text); + let styled = self.bold.apply(&decoded); + self.italic.apply(&styled.to_string()).to_string() + } + + fn strikethrough(&self, text: &str) -> String { + self.strikethrough + .apply(&decode_html_entities(text)) + .to_string() + } + + fn underline(&self, text: &str) -> String { + format!("\x1b[4m{}\x1b[24m", decode_html_entities(text)) + } + + fn code(&self, text: &str) -> String { + self.code.apply(text).to_string() + } + + fn link(&self, text: &str, url: &str) -> String { + let mut result = String::new(); + result.push_str("\x1b]8;;"); + result.push_str(url); + result.push_str("\x1b\\"); + result.push_str(&self.link.apply(&decode_html_entities(text)).to_string()); + result.push_str("\x1b]8;;\x1b\\"); + result.push(' '); + result.push_str(&self.link_url.apply(&format!("({})", url)).to_string()); + result + } + + fn image(&self, alt: &str, _url: &str) -> String { + format!("[🖼 {}]", alt) + } + + fn footnote(&self, text: &str) -> String { + text.to_string() + } + + fn dimmed(&self, text: &str) -> String { + Style::new().dimmed().apply(text).to_string() + } +} + +impl HeadingStyler for Theme { + fn h1(&self, text: &str) -> String { + self.heading1.apply(text).to_string() + } + + fn h2(&self, text: &str) -> String { + self.heading2.apply(text).to_string() + } + + fn h3(&self, text: &str) -> String { + self.heading3.apply(text).to_string() + } + + fn h4(&self, text: &str) -> String { + self.heading4.apply(text).to_string() + } + + fn h5(&self, text: &str) -> String { + self.heading5.apply(text).to_string() + } + + fn h6(&self, text: &str) -> String { + self.heading6.apply(text).to_string() + } +} + +impl ListStyler for Theme { + fn bullet_dash(&self, text: &str) -> String { + self.bullet_dash.apply(text).to_string() + } + + fn bullet_asterisk(&self, text: &str) -> String { + self.bullet_asterisk.apply(text).to_string() + } + + fn bullet_plus(&self, text: &str) -> String { + self.bullet_plus.apply(text).to_string() + } + + fn bullet_plus_expand(&self, text: &str) -> String { + self.bullet_plus_expand.apply(text).to_string() + } + + fn number(&self, text: &str) -> String { + self.list_number.apply(text).to_string() + } + + fn checkbox_checked(&self, text: &str) -> String { + self.checkbox_checked.apply(text).to_string() + } + + fn checkbox_unchecked(&self, text: &str) -> String { + self.checkbox_unchecked.apply(text).to_string() + } +} + +impl TableStyler for Theme { + fn border(&self, text: &str) -> String { + self.table_border.apply(text).to_string() + } + + fn header(&self, text: &str) -> String { + self.table_header.apply(text).to_string() + } +} + +impl Theme { + /// Detects the terminal theme (dark or light) and returns the appropriate + /// theme. + pub fn detect() -> Self { + use crate::utils::{ThemeMode, detect_theme_mode}; + + match detect_theme_mode() { + ThemeMode::Light => Self::light(), + ThemeMode::Dark => Self::dark(), + } + } + + /// Dark theme (default). + pub fn dark() -> Self { + Self { + // Inline + bold: Style::new().bold(), + italic: Style::new().italic(), + code: Style::new().fg(Color::Yellow), + strikethrough: Style::new().strikethrough().dimmed(), + link: Style::new().fg(Color::Cyan).underline(), + link_url: Style::new().fg(Color::Blue).dimmed(), + + // Headings + heading1: Style::new().fg(Color::Magenta).bold(), + heading2: Style::new().fg(Color::Blue).bold(), + heading3: Style::new().fg(Color::Cyan).bold(), + heading4: Style::new().fg(Color::Green).bold(), + heading5: Style::new().fg(Color::Yellow).bold(), + heading6: Style::new().fg(Color::White).bold(), + + // Lists + bullet_dash: Style::new().fg(Color::Cyan), + bullet_asterisk: Style::new().fg(Color::Green), + bullet_plus: Style::new().fg(Color::Yellow), + bullet_plus_expand: Style::new().fg(Color::Magenta), + list_number: Style::new().fg(Color::Cyan), + checkbox_checked: Style::new().fg(Color::Green), + checkbox_unchecked: Style::new().fg(Color::Red), + + // Tables + table_header: Style::new().bold(), + table_border: Style::new().fg(Color::BrightBlack), + table_cell: Style::new(), + + // Quotes + blockquote: Style::new().italic().dimmed(), + blockquote_border: Style::new().fg(Color::BrightBlack), + think: Style::new().italic().fg(Color::BrightBlack), + think_border: Style::new().fg(Color::BrightBlack), + + // Code block + code_block_lang: Style::new().fg(Color::BrightBlack).italic(), + + // HR + hr: Style::new().fg(Color::BrightBlack), + } + } + + /// Light theme for light terminal backgrounds. + pub fn light() -> Self { + Self { + // Inline + bold: Style::new().bold(), + italic: Style::new().italic(), + code: Style::new().fg(Color::Red), + strikethrough: Style::new().strikethrough().dimmed(), + link: Style::new().fg(Color::Blue).underline(), + link_url: Style::new().fg(Color::Cyan).dimmed(), + + // Headings + heading1: Style::new().fg(Color::Magenta).bold(), + heading2: Style::new().fg(Color::Blue).bold(), + heading3: Style::new().fg(Color::Cyan).bold(), + heading4: Style::new().fg(Color::Green).bold(), + heading5: Style::new().fg(Color::Yellow).bold(), + heading6: Style::new().fg(Color::Black).bold(), + + // Lists + bullet_dash: Style::new().fg(Color::Blue), + bullet_asterisk: Style::new().fg(Color::Green), + bullet_plus: Style::new().fg(Color::Magenta), + bullet_plus_expand: Style::new().fg(Color::Cyan), + list_number: Style::new().fg(Color::Blue), + checkbox_checked: Style::new().fg(Color::Green), + checkbox_unchecked: Style::new().fg(Color::Red), + + // Tables + table_header: Style::new().bold(), + table_border: Style::new().fg(Color::Black), + table_cell: Style::new(), + + // Quotes + blockquote: Style::new().italic().dimmed(), + blockquote_border: Style::new().fg(Color::Black), + think: Style::new().italic().fg(Color::Black), + think_border: Style::new().fg(Color::Black), + + // Code block + code_block_lang: Style::new().fg(Color::Black).italic(), + + // HR + hr: Style::new().fg(Color::Black), + } + } +} + +/// Test styler that outputs readable HTML-like tags. +#[cfg(test)] +pub struct TagStyler; + +#[cfg(test)] +impl InlineStyler for TagStyler { + fn text(&self, text: &str) -> String { + decode_html_entities(text) + } + + fn bold(&self, text: &str) -> String { + format!("{}", decode_html_entities(text)) + } + + fn italic(&self, text: &str) -> String { + format!("{}", decode_html_entities(text)) + } + + fn bold_italic(&self, text: &str) -> String { + format!("{}", decode_html_entities(text)) + } + + fn strikethrough(&self, text: &str) -> String { + format!("{}", decode_html_entities(text)) + } + + fn underline(&self, text: &str) -> String { + format!("{}", decode_html_entities(text)) + } + + fn code(&self, text: &str) -> String { + format!("{}", text) + } + + fn link(&self, text: &str, url: &str) -> String { + format!("{}", url, decode_html_entities(text)) + } + + fn image(&self, alt: &str, url: &str) -> String { + format!("\"{}\"", alt, url) + } + + fn footnote(&self, text: &str) -> String { + format!("{}", text) + } + + fn dimmed(&self, text: &str) -> String { + format!("{}", text) + } +} + +#[cfg(test)] +impl HeadingStyler for TagStyler { + fn h1(&self, text: &str) -> String { + format!("

{}

", text) + } + + fn h2(&self, text: &str) -> String { + format!("

{}

", text) + } + + fn h3(&self, text: &str) -> String { + format!("

{}

", text) + } + + fn h4(&self, text: &str) -> String { + format!("

{}

", text) + } + + fn h5(&self, text: &str) -> String { + format!("
{}
", text) + } + + fn h6(&self, text: &str) -> String { + format!("
{}
", text) + } +} + +#[cfg(test)] +impl ListStyler for TagStyler { + fn bullet_dash(&self, text: &str) -> String { + format!("{}", text) + } + + fn bullet_asterisk(&self, text: &str) -> String { + format!("{}", text) + } + + fn bullet_plus(&self, text: &str) -> String { + format!("{}", text) + } + + fn bullet_plus_expand(&self, text: &str) -> String { + format!("{}", text) + } + + fn number(&self, text: &str) -> String { + format!("{}", text) + } + + fn checkbox_checked(&self, text: &str) -> String { + format!("{}", text) + } + + fn checkbox_unchecked(&self, text: &str) -> String { + format!("{}", text) + } +} + +#[cfg(test)] +impl TableStyler for TagStyler { + fn border(&self, text: &str) -> String { + Theme::default().border(text) + } + + fn header(&self, text: &str) -> String { + Theme::default().header(text) + } +} diff --git a/crates/forge_markdown_stream/src/utils.rs b/crates/forge_markdown_stream/src/utils.rs new file mode 100644 index 0000000000000000000000000000000000000000..6145ef7360b4c6f30cd79bc7c870963ab4591ef4 --- /dev/null +++ b/crates/forge_markdown_stream/src/utils.rs @@ -0,0 +1,550 @@ +//! Utility functions for the markdown renderer. + +use std::sync::OnceLock; +use std::time::Duration; + +use streamdown_ansi::utils::{extract_ansi_codes, parse_sgr_params, visible, visible_length}; +use unicode_segmentation::UnicodeSegmentation; +use unicode_width::UnicodeWidthStr; + +/// Terminal theme mode (dark or light). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ThemeMode { + /// Dark terminal background. + Dark, + /// Light terminal background. + Light, +} + +/// Maximum time to wait for a terminal color query response. +const THEME_DETECT_TIMEOUT: Duration = Duration::from_millis(100); + +/// Process-wide cache for the detected terminal theme mode. +static THEME_MODE: OnceLock = OnceLock::new(); + +/// Detects the terminal theme mode (dark or light), querying the terminal at +/// most once per process lifetime. Subsequent calls return the cached result. +/// Falls back to dark mode if the terminal does not respond within the timeout. +pub fn detect_theme_mode() -> ThemeMode { + *THEME_MODE.get_or_init(|| { + use terminal_colorsaurus::{QueryOptions, ThemeMode as ColorsaurusThemeMode, theme_mode}; + + let mut opts = QueryOptions::default(); + opts.timeout = THEME_DETECT_TIMEOUT; + match theme_mode(opts) { + Ok(ColorsaurusThemeMode::Light) => ThemeMode::Light, + Ok(ColorsaurusThemeMode::Dark) | Err(_) => ThemeMode::Dark, + } + }) +} + +#[derive(Debug, Clone, PartialEq, Eq)] +enum WrapAtom<'a> { + Escape(&'a str), + Grapheme(&'a str), +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct WrapSegment { + separator: String, + word: String, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct WrapLayout<'a> { + next_width: usize, + first_prefix: &'a str, + next_prefix: &'a str, +} + +/// Wraps ANSI-styled text while preserving explicit whitespace between words. +/// +/// Unlike the upstream streamdown wrapper, this keeps the original separator +/// string between tokens instead of reconstructing it from CJK heuristics. +pub(crate) fn wrap_text_preserving_spaces( + text: &str, + first_width: usize, + next_width: usize, + first_prefix: &str, + next_prefix: &str, +) -> Vec { + if first_width == 0 && next_width == 0 { + return Vec::new(); + } + + let segments = wrap_segments(text); + if segments.is_empty() { + return Vec::new(); + } + + let mut lines = Vec::new(); + let mut current_line = String::new(); + let mut current_style: Vec = Vec::new(); + let mut current_width = first_width; + let layout = WrapLayout { next_width, first_prefix, next_prefix }; + + for segment in segments { + let line_width = visible_length(¤t_line); + let separator = if current_line.is_empty() { + "" + } else { + segment.separator.as_str() + }; + let combined_width = visible_length(separator) + visible_length(&segment.word); + + if !current_line.is_empty() && line_width + combined_width <= current_width { + current_line.push_str(separator); + apply_style_transition(&mut current_style, separator); + current_line.push_str(&segment.word); + apply_style_transition(&mut current_style, &segment.word); + continue; + } + + if current_line.is_empty() && visible_length(&segment.word) <= current_width { + current_line.push_str(&segment.word); + apply_style_transition(&mut current_style, &segment.word); + continue; + } + + if !current_line.is_empty() { + push_wrapped_line(&mut lines, ¤t_line, first_prefix, next_prefix); + current_line = current_style.join(""); + current_width = next_width; + } + + append_wrapped_word( + &mut lines, + &mut current_line, + &mut current_style, + &segment.word, + &mut current_width, + layout, + ); + } + + push_wrapped_line(&mut lines, ¤t_line, first_prefix, next_prefix); + lines +} + +/// Wraps ANSI-styled inline text without prefixes while preserving explicit +/// spaces. +pub(crate) fn simple_wrap_preserving_spaces(text: &str, width: usize) -> Vec { + if width == 0 || text.is_empty() { + return vec![text.to_string()]; + } + + let lines = wrap_text_preserving_spaces(text, width, width, "", ""); + if lines.is_empty() { + vec![String::new()] + } else { + lines + } +} + +fn append_wrapped_word( + lines: &mut Vec, + current_line: &mut String, + current_style: &mut Vec, + word: &str, + current_width: &mut usize, + layout: WrapLayout<'_>, +) { + let mut remainder = word; + + while !remainder.is_empty() { + let line_width = visible_length(current_line); + let mut available = current_width.saturating_sub(line_width); + + if available == 0 { + push_wrapped_line(lines, current_line, layout.first_prefix, layout.next_prefix); + *current_line = current_style.join(""); + *current_width = layout.next_width; + available = (*current_width).max(1); + } + + if visible_length(remainder) <= available { + current_line.push_str(remainder); + apply_style_transition(current_style, remainder); + break; + } + + let prefix = take_prefix_fitting(remainder, available) + .or_else(|| take_prefix_fitting(remainder, 1)) + .unwrap_or(remainder); + + current_line.push_str(prefix); + apply_style_transition(current_style, prefix); + remainder = remainder.strip_prefix(prefix).unwrap_or_default(); + + if !remainder.is_empty() { + push_wrapped_line(lines, current_line, layout.first_prefix, layout.next_prefix); + *current_line = current_style.join(""); + *current_width = layout.next_width; + } + } +} + +fn take_prefix_fitting(text: &str, max_width: usize) -> Option<&str> { + if text.is_empty() { + return None; + } + + let mut width = 0; + let mut prefix_end = 0; + let mut consumed_visible = false; + + for atom in parse_atoms(text) { + match atom { + WrapAtom::Escape(sequence) => prefix_end += sequence.len(), + WrapAtom::Grapheme(grapheme) => { + let grapheme_width = UnicodeWidthStr::width(grapheme); + if consumed_visible && width + grapheme_width > max_width { + break; + } + if !consumed_visible && grapheme_width > max_width { + prefix_end += grapheme.len(); + break; + } + + prefix_end += grapheme.len(); + width += grapheme_width; + consumed_visible = true; + } + } + } + + (prefix_end > 0).then(|| text.get(..prefix_end)).flatten() +} + +fn wrap_segments(text: &str) -> Vec { + let mut segments = Vec::new(); + let mut current = String::new(); + let mut current_is_whitespace = None; + let mut separator = String::new(); + + for atom in parse_atoms(text) { + match atom { + WrapAtom::Escape(sequence) => current.push_str(sequence), + WrapAtom::Grapheme(grapheme) => { + let is_whitespace = grapheme.chars().all(char::is_whitespace); + match current_is_whitespace { + Some(kind) if kind != is_whitespace => { + if kind { + separator.push_str(¤t); + current.clear(); + } else { + segments.push(WrapSegment { + separator: std::mem::take(&mut separator), + word: std::mem::take(&mut current), + }); + } + current_is_whitespace = Some(is_whitespace); + } + None => current_is_whitespace = Some(is_whitespace), + _ => {} + } + + current.push_str(grapheme); + } + } + } + + match (current_is_whitespace, current.is_empty()) { + (Some(true), false) => separator.push_str(¤t), + (Some(false), false) => segments.push(WrapSegment { separator, word: current }), + _ => {} + } + + segments +} + +fn parse_atoms(text: &str) -> Vec> { + let mut atoms = Vec::new(); + let bytes = text.as_bytes(); + let mut index = 0; + + while index < bytes.len() { + let Some(current_byte) = bytes.get(index) else { + break; + }; + + if *current_byte != 0x1b { + let next_escape = text + .get(index..) + .and_then(|s| s.find('\x1b')) + .map(|offset| index + offset) + .unwrap_or(text.len()); + + if let Some(slice) = text.get(index..next_escape) { + for grapheme in slice.graphemes(true) { + atoms.push(WrapAtom::Grapheme(grapheme)); + } + } + + index = next_escape; + continue; + } + + let end = match bytes.get(index + 1) { + Some(b'[') => parse_csi_escape(bytes, index), + Some(b']') => parse_osc_escape(bytes, index), + Some(_) => (index + 2).min(bytes.len()), + None => bytes.len(), + }; + + if let Some(sequence) = text.get(index..end) { + atoms.push(WrapAtom::Escape(sequence)); + } + + index = end; + } + + atoms +} + +fn parse_csi_escape(bytes: &[u8], start: usize) -> usize { + let mut index = start + 2; + while let Some(byte) = bytes.get(index) { + if (0x40..=0x7e).contains(byte) { + return index + 1; + } + + index += 1; + } + + bytes.len() +} + +fn parse_osc_escape(bytes: &[u8], start: usize) -> usize { + let mut index = start + 2; + while let Some(byte) = bytes.get(index) { + if *byte == 0x07 { + return index + 1; + } + + if *byte == 0x1b && bytes.get(index + 1) == Some(&b'\\') { + return index + 2; + } + + index += 1; + } + + bytes.len() +} + +fn apply_style_transition(current_style: &mut Vec, text: &str) { + current_style.extend(extract_ansi_codes(text)); + *current_style = collapse_ansi_codes(current_style); +} + +fn collapse_ansi_codes(code_list: &[String]) -> Vec { + let mut bold = false; + let mut italic = false; + let mut underline = false; + let mut strikeout = false; + let mut dim = false; + let mut fg_color: Option = None; + let mut bg_color: Option = None; + + for code in code_list { + let params = parse_sgr_params(code); + let mut index = 0; + + while let Some(¶m) = params.get(index) { + match param { + 0 => { + bold = false; + italic = false; + underline = false; + strikeout = false; + dim = false; + fg_color = None; + bg_color = None; + } + 1 => bold = true, + 2 => dim = true, + 3 => italic = true, + 4 => underline = true, + 9 => strikeout = true, + 22 => { + bold = false; + dim = false; + } + 23 => italic = false, + 24 => underline = false, + 29 => strikeout = false, + 30..=37 | 90..=97 => fg_color = Some(format!("\x1b[{param}m")), + 39 => fg_color = None, + 40..=47 | 100..=107 => bg_color = Some(format!("\x1b[{param}m")), + 49 => bg_color = None, + 38 => { + if let Some([2, red, green, blue]) = params.get(index + 1..index + 5) { + fg_color = Some(format!("\x1b[38;2;{red};{green};{blue}m")); + index += 4; + } + } + 48 => { + if let Some([2, red, green, blue]) = params.get(index + 1..index + 5) { + bg_color = Some(format!("\x1b[48;2;{red};{green};{blue}m")); + index += 4; + } + } + _ => {} + } + + index += 1; + } + } + + let mut result = Vec::new(); + let mut sgr_parts = Vec::new(); + + if bold { + sgr_parts.push("1"); + } + if dim { + sgr_parts.push("2"); + } + if italic { + sgr_parts.push("3"); + } + if underline { + sgr_parts.push("4"); + } + if strikeout { + sgr_parts.push("9"); + } + + if !sgr_parts.is_empty() { + result.push(format!("\x1b[{}m", sgr_parts.join(";"))); + } + if let Some(fg_color) = fg_color { + result.push(fg_color); + } + if let Some(bg_color) = bg_color { + result.push(bg_color); + } + + result +} + +fn push_wrapped_line( + lines: &mut Vec, + current_line: &str, + first_prefix: &str, + next_prefix: &str, +) { + if visible(current_line).trim().is_empty() { + return; + } + + let prefix = if lines.is_empty() { + first_prefix + } else { + next_prefix + }; + lines.push(format!("{prefix}{current_line}")); +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + use streamdown_ansi::utils::visible; + + use super::{simple_wrap_preserving_spaces, wrap_text_preserving_spaces}; + + #[test] + fn test_simple_wrap_preserving_spaces_keeps_korean_word_boundaries() { + let fixture = "한글 공백 보존 문장"; + let actual = simple_wrap_preserving_spaces(fixture, 8); + let expected = vec![ + "한글".to_string(), + "공백".to_string(), + "보존".to_string(), + "문장".to_string(), + ]; + + assert_eq!(actual, expected); + } + + #[test] + fn test_simple_wrap_preserving_spaces_splits_long_tokens() { + let fixture = "supercalifragilistic"; + let actual = simple_wrap_preserving_spaces(fixture, 5); + let expected = vec![ + "super".to_string(), + "calif".to_string(), + "ragil".to_string(), + "istic".to_string(), + ]; + + assert_eq!(actual, expected); + } + + #[test] + fn test_wrap_text_preserving_spaces_keeps_multiple_spaces_on_same_line() { + let fixture = "한글 공백 보존"; + let actual = wrap_text_preserving_spaces(fixture, 40, 40, "", ""); + let expected = vec!["한글 공백 보존".to_string()]; + + assert_eq!(actual, expected); + } + + #[test] + fn test_wrap_text_preserving_spaces_applies_prefixes_after_wrap() { + let fixture = "한글 공백 검증"; + let actual = wrap_text_preserving_spaces(fixture, 4, 4, "> ", " "); + let expected = vec![ + "> 한글".to_string(), + " 공백".to_string(), + " 검증".to_string(), + ]; + + assert_eq!(actual, expected); + } + + #[test] + fn test_wrap_text_preserving_spaces_preserves_link_separator_after_osc_escape() { + let fixture = concat!( + "\x1b]8;;https://example.com\x1b\\", + "link", + "\x1b]8;;\x1b\\", + " ", + "\x1b[34m(https://x.co)\x1b[39m" + ); + let actual = wrap_text_preserving_spaces(fixture, 4, 14, "", "") + .into_iter() + .map(|line| visible(&line)) + .collect::>(); + let expected = vec!["link".to_string(), "(https://x.co)".to_string()]; + + assert_eq!(actual, expected); + } + + #[test] + fn test_simple_wrap_preserving_spaces_keeps_grapheme_clusters_intact() { + let fixture = "👨‍👩‍👧‍👦 a\u{0301} 한글"; + let actual = simple_wrap_preserving_spaces(fixture, 2); + let expected = vec![ + "👨‍👩‍👧‍👦".to_string(), + "a\u{0301}".to_string(), + "한".to_string(), + "글".to_string(), + ]; + + assert_eq!(actual, expected); + } + + #[test] + fn test_wrap_text_preserving_spaces_reapplies_ansi_style_after_wrap() { + let fixture = "\x1b[31mabcdef\x1b[39m"; + let actual = wrap_text_preserving_spaces(fixture, 3, 3, "", ""); + let expected_visible = vec!["abc".to_string(), "def".to_string()]; + let actual_visible = actual.iter().map(|line| visible(line)).collect::>(); + + assert_eq!(actual_visible, expected_visible); + assert!(actual[0].contains("\x1b[31m")); + assert!(actual[1].contains("\x1b[31m")); + assert!(actual[1].ends_with("\x1b[39m")); + } +} diff --git a/crates/forge_snaps/src/lib.rs b/crates/forge_snaps/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..eebfa61391126a5025006df23c88f15744b9be69 --- /dev/null +++ b/crates/forge_snaps/src/lib.rs @@ -0,0 +1,5 @@ +// Export the modules +mod service; + +// Re-export the SnapshotInfo struct and SnapshotId +pub use service::*; diff --git a/crates/forge_snaps/src/service.rs b/crates/forge_snaps/src/service.rs new file mode 100644 index 0000000000000000000000000000000000000000..4fa11b53dc07d409f5bc2ca80117dc7949493890 --- /dev/null +++ b/crates/forge_snaps/src/service.rs @@ -0,0 +1,253 @@ +use std::path::PathBuf; + +use anyhow::{Context, Result}; +use forge_domain::Snapshot; +use forge_fs::ForgeFS; + +/// Implementation of the SnapshotService +#[derive(Debug)] +pub struct SnapshotService { + /// Base directory for storing snapshots + snapshots_directory: PathBuf, +} + +impl SnapshotService { + /// Create a new FileSystemSnapshotService with a specific home path + pub fn new(snapshot_base_dir: PathBuf) -> Self { + Self { snapshots_directory: snapshot_base_dir } + } +} + +impl SnapshotService { + pub async fn create_snapshot(&self, path: PathBuf) -> Result { + let snapshot = Snapshot::create(path)?; + + // Create intermediary directories if they don't exist + let snapshot_path = snapshot.snapshot_path(Some(self.snapshots_directory.clone())); + if let Some(parent) = PathBuf::from(&snapshot_path).parent() { + ForgeFS::create_dir_all(parent).await?; + } + + let content = ForgeFS::read(&snapshot.path).await?; + let path = snapshot.snapshot_path(Some(self.snapshots_directory.clone())); + ForgeFS::write(path, content).await?; + Ok(snapshot) + } + + /// Find the most recent snapshot for a given path based on filename + /// timestamp + async fn find_recent_snapshot(snapshot_dir: &PathBuf) -> Result> { + let mut latest_path = None; + let mut latest_filename = None; + let mut dir = ForgeFS::read_dir(&snapshot_dir).await?; + + while let Some(entry) = dir.next_entry().await? { + let filename = entry.file_name().to_string_lossy().to_string(); + if filename.ends_with(".snap") + && (latest_filename.is_none() || filename > latest_filename.clone().unwrap()) + { + latest_filename = Some(filename); + latest_path = Some(entry.path()); + } + } + + Ok(latest_path) + } + + pub async fn undo_snapshot(&self, path: PathBuf) -> Result<()> { + let snapshot = Snapshot::create(path.clone())?; + + // All the snaps for `path` are stored in `snapshot.path_hash()` directory. + let snapshot_dir = self.snapshots_directory.join(snapshot.path_hash()); + + // Check if the `snapshot_dir` exists + if !ForgeFS::exists(&snapshot_dir) { + return Err(anyhow::anyhow!("No snapshots found for {path:?}")); + } + + // Retrieve the latest snapshot path + let snapshot_path = Self::find_recent_snapshot(&snapshot_dir) + .await? + .context(format!("No valid snapshots found for {path:?}"))?; + + // Restore the content + let content = ForgeFS::read(&snapshot_path).await?; + ForgeFS::write(&path, content).await?; + + // Remove the used snapshot + ForgeFS::remove_file(&snapshot_path).await?; + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use tempfile::TempDir; + + use super::*; + + // Test helpers + struct TestContext { + _temp_dir: TempDir, + _snapshots_dir: PathBuf, + test_file: PathBuf, + service: SnapshotService, + } + + impl TestContext { + async fn new() -> Result { + let temp_dir = TempDir::new()?; + let snapshots_dir = temp_dir.path().join("snapshots"); + // Canonicalize the temp directory path to ensure consistency + let temp_path = temp_dir + .path() + .canonicalize() + .unwrap_or_else(|_| temp_dir.path().to_path_buf()); + let test_file = temp_path.join("test.txt"); + let service = SnapshotService::new(snapshots_dir.clone()); + + Ok(Self { + _temp_dir: temp_dir, + _snapshots_dir: snapshots_dir, + test_file, + service, + }) + } + + async fn write_content(&self, content: &str) -> Result<()> { + ForgeFS::write(&self.test_file, content.as_bytes()).await + } + + async fn read_content(&self) -> Result { + let content = ForgeFS::read(&self.test_file).await?; + Ok(String::from_utf8(content)?) + } + + async fn create_snapshot(&self) -> Result { + self.service.create_snapshot(self.test_file.clone()).await + } + + async fn undo_snapshot(&self) -> Result<()> { + self.service.undo_snapshot(self.test_file.clone()).await + } + } + + #[tokio::test] + async fn test_create_snapshot() -> Result<()> { + // Arrange + let ctx = TestContext::new().await?; + let test_content = "Hello, World!"; + + // Act + ctx.write_content(test_content).await?; + let snapshot = ctx.create_snapshot().await?; + + // Assert + let snapshot_content = ForgeFS::read(&snapshot.path).await?; + assert_eq!(String::from_utf8(snapshot_content)?, test_content); + + Ok(()) + } + + #[tokio::test] + async fn test_undo_snapshot() -> Result<()> { + // Arrange + let ctx = TestContext::new().await?; + let initial_content = "Initial content"; + let modified_content = "Modified content"; + + // Act + ctx.write_content(initial_content).await?; + ctx.create_snapshot().await?; + ctx.write_content(modified_content).await?; + ctx.undo_snapshot().await?; + + // Assert + assert_eq!(ctx.read_content().await?, initial_content); + + Ok(()) + } + + #[tokio::test] + async fn test_undo_snapshot_no_snapshots() -> Result<()> { + // Arrange + let ctx = TestContext::new().await?; + + // Act + ctx.write_content("test content").await?; + let result = ctx.undo_snapshot().await; + + // Assert + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .contains("No snapshots found") + ); + + Ok(()) + } + + #[tokio::test] + async fn test_undo_snapshot_after_file_deletion() -> Result<()> { + // Arrange + let ctx = TestContext::new().await?; + let initial_content = "Initial content"; + + // Act + ctx.write_content(initial_content).await?; + ctx.create_snapshot().await?; + ForgeFS::remove_file(&ctx.test_file).await?; + + // Assert - undo should succeed and recreate the file from snapshot + ctx.undo_snapshot().await?; + assert_eq!(ctx.read_content().await?, initial_content); + + Ok(()) + } + + #[tokio::test] + async fn test_multiple_snapshots() -> Result<()> { + // Arrange + let ctx = TestContext::new().await?; + + // Act + ctx.write_content("Initial content").await?; + ctx.create_snapshot().await?; + + ctx.write_content("Second content").await?; + ctx.create_snapshot().await?; + + ctx.write_content("Final content").await?; + ctx.undo_snapshot().await?; + + // Assert + assert_eq!(ctx.read_content().await?, "Second content"); + + Ok(()) + } + + #[tokio::test] + async fn test_multiple_snapshots_undo_twice() -> Result<()> { + // Arrange + let ctx = TestContext::new().await?; + + // Act + ctx.write_content("Initial content").await?; + ctx.create_snapshot().await?; + + ctx.write_content("Second content").await?; + ctx.create_snapshot().await?; + + ctx.write_content("Final content").await?; + ctx.undo_snapshot().await?; + ctx.undo_snapshot().await?; + + // Assert + assert_eq!(ctx.read_content().await?, "Initial content"); + + Ok(()) + } +}