From 603b35b45184c0888ac74327ef4f8b7365828011 Mon Sep 17 00:00:00 2001 From: blankll Date: Fri, 31 Jul 2026 00:21:01 +0800 Subject: [PATCH 1/8] refactor: unify capability types with data-studio-agent - Remove duplicate capabilities/types.rs and capabilities/registry.rs - Import RiskLevel, SourceKind, CapabilityHandler, Capability, CapabilityRegistry from data-studio-agent instead - Replace SourceKind::DocKit with SourceKind::AppLocal - Update init_registry() call to pass registration function array - Update agent/executor.rs and agent_adapters.rs to use data-studio-agent paths --- src-tauri/Cargo.lock | 193 +++++++++-------- src-tauri/Cargo.toml | 3 +- src-tauri/src/agent/executor.rs | 2 +- src-tauri/src/agent_adapters.rs | 2 +- src-tauri/src/capabilities/commands.rs | 60 ++++-- src-tauri/src/capabilities/dockit.rs | 6 +- src-tauri/src/capabilities/dynamo.rs | 4 +- src-tauri/src/capabilities/es.rs | 6 +- src-tauri/src/capabilities/mod.rs | 5 - src-tauri/src/capabilities/mongo.rs | 4 +- src-tauri/src/capabilities/registry.rs | 278 ------------------------- src-tauri/src/capabilities/types.rs | 71 ------- src-tauri/src/lib.rs | 7 +- src-tauri/src/tests.rs | 18 +- 14 files changed, 166 insertions(+), 493 deletions(-) delete mode 100644 src-tauri/src/capabilities/registry.rs delete mode 100644 src-tauri/src/capabilities/types.rs diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index 93b353c5..ca4d5a80 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -41,9 +41,9 @@ dependencies = [ [[package]] name = "aes" -version = "0.9.1" +version = "0.9.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f1fc76eaeac4c9164506c466d4ffdd8ec9d0c5bf57ee97177c4d8eceb3a0e138" +checksum = "f8eb277bec05f56a0e0591f155a484cbd0f4f07ff2905051a48c72f004f7ed58" dependencies = [ "cipher 0.5.2", "cpubits", @@ -71,7 +71,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e22c0c90bbe8d4f77c3ca9ddabe41a1f8382d6fc1f7cea89459d0f320371f972" dependencies = [ "aead 0.6.0-rc.10", - "aes 0.9.1", + "aes 0.9.2", "cipher 0.5.2", "ctr 0.10.1", "ghash 0.6.0", @@ -384,9 +384,9 @@ checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" [[package]] name = "aws-config" -version = "1.10.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "701418aa459dac33e50a0f8e818e5662a16bc018a6ac7423659b70f3799d67a8" +checksum = "1b180a3c8b55960db3426d8964b8745e652466a1a49fe1a2eda828046d30b5e4" dependencies = [ "aws-credential-types", "aws-runtime", @@ -404,7 +404,7 @@ dependencies = [ "bytes", "fastrand", "hex", - "http 1.4.2", + "http 1.5.0", "sha1 0.10.7", "time", "tokio", @@ -451,9 +451,9 @@ dependencies = [ [[package]] name = "aws-runtime" -version = "1.9.0" +version = "1.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a6b50a43f3ccdf331521c6d6c68b7cc9668b6e09d439ebda9569df5722324d76" +checksum = "c9007227e10b5fed2f3e0a2beff489211e2b5604c400b7a9d5d81ca9d64c24bb" dependencies = [ "aws-credential-types", "aws-sigv4", @@ -466,7 +466,7 @@ dependencies = [ "bytes", "bytes-utils", "fastrand", - "http 1.4.2", + "http 1.5.0", "http-body 1.1.0", "percent-encoding", "pin-project-lite", @@ -476,9 +476,9 @@ dependencies = [ [[package]] name = "aws-sdk-cloudwatch" -version = "1.121.0" +version = "1.123.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1eb29868692a5aeb06a33dfa753b2b47fb94971204a15ba8ceec7c21a7bcf621" +checksum = "72d33759772cc816ccd6d6c10a58f0c193405b6ed60e89ed3f5e3ea7f29f97d2" dependencies = [ "arc-swap", "aws-credential-types", @@ -497,7 +497,7 @@ dependencies = [ "fastrand", "flate2", "http 0.2.12", - "http 1.4.2", + "http 1.5.0", "http-body 1.1.0", "http-body-util", "regex-lite", @@ -506,9 +506,9 @@ dependencies = [ [[package]] name = "aws-sdk-dynamodb" -version = "1.118.0" +version = "1.119.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d1d35ca493e71083520b9d3492c205cc356aea190195107a80dce07633d5b027" +checksum = "b0176c1927f2e065f43cecde84602f769959b34ab53a317da92a64022ba09345" dependencies = [ "arc-swap", "aws-credential-types", @@ -525,16 +525,16 @@ dependencies = [ "bytes", "fastrand", "http 0.2.12", - "http 1.4.2", + "http 1.5.0", "regex-lite", "tracing", ] [[package]] name = "aws-sdk-sso" -version = "1.104.0" +version = "1.105.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b53416d16c278234845392e38d93bd4481d2f09daa0f005a2277f0aa91f59c22" +checksum = "6ffd0fbe7873cb548a7aa60f9573c268fff94155397fd4f14dc9f1ecaaab8516" dependencies = [ "arc-swap", "aws-credential-types", @@ -551,16 +551,16 @@ dependencies = [ "bytes", "fastrand", "http 0.2.12", - "http 1.4.2", + "http 1.5.0", "regex-lite", "tracing", ] [[package]] name = "aws-sdk-ssooidc" -version = "1.106.0" +version = "1.107.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cc9b706c3305ed0285d5b1b696c747aa34950f830fb03e3e6c76890f99b9f188" +checksum = "175763eb222a46377df7aa257a3bca980ab3e96703fefc8f4d0b8da6ad2e254c" dependencies = [ "arc-swap", "aws-credential-types", @@ -577,16 +577,16 @@ dependencies = [ "bytes", "fastrand", "http 0.2.12", - "http 1.4.2", + "http 1.5.0", "regex-lite", "tracing", ] [[package]] name = "aws-sdk-sts" -version = "1.109.0" +version = "1.110.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32d214cdfa5bbe17f117e76a7643fadf32a5234fb597322ef8b1fb4b2f17dbbd" +checksum = "dd8b14781dfbff48984017d57167b6ea0b6471c6920ec52b44a2677c7feb3c13" dependencies = [ "arc-swap", "aws-credential-types", @@ -604,7 +604,7 @@ dependencies = [ "aws-types", "fastrand", "http 0.2.12", - "http 1.4.2", + "http 1.5.0", "regex-lite", "tracing", ] @@ -624,7 +624,7 @@ dependencies = [ "hex", "hmac 0.13.0", "http 0.2.12", - "http 1.4.2", + "http 1.5.0", "percent-encoding", "sha2 0.11.0", "time", @@ -666,7 +666,7 @@ dependencies = [ "bytes", "flate2", "futures-util", - "http 1.4.2", + "http 1.5.0", "http-body 1.1.0", "http-body-util", "pin-project-lite", @@ -685,7 +685,7 @@ dependencies = [ "bytes-utils", "futures-core", "futures-util", - "http 1.4.2", + "http 1.5.0", "http-body 1.1.0", "http-body-util", "percent-encoding", @@ -704,7 +704,7 @@ dependencies = [ "aws-smithy-runtime-api", "aws-smithy-types", "h2", - "http 1.4.2", + "http 1.5.0", "hyper", "hyper-rustls", "hyper-util", @@ -753,9 +753,9 @@ dependencies = [ [[package]] name = "aws-smithy-runtime" -version = "1.12.0" +version = "1.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bea94a9ff8464016338c851e24b472d7131c388c88898a502e781815b2ee6045" +checksum = "07505b34e8f4b3591a4fa69e9792b52289b95488dbbc68c3c0075b7bedb245e1" dependencies = [ "aws-smithy-async", "aws-smithy-http", @@ -767,7 +767,7 @@ dependencies = [ "bytes", "fastrand", "http 0.2.12", - "http 1.4.2", + "http 1.5.0", "http-body 0.4.6", "http-body 1.1.0", "http-body-util", @@ -779,16 +779,16 @@ dependencies = [ [[package]] name = "aws-smithy-runtime-api" -version = "1.13.0" +version = "1.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22ed1ebe6e0a95ea84570225f5a8208dec4b8f77e61a9b0d6f51773fcb4612f0" +checksum = "3b98f2e1fd67ec06618f9c291e5e495a468e60519e44c9c1979cd0521f3affdb" dependencies = [ "aws-smithy-async", "aws-smithy-runtime-api-macros", "aws-smithy-types", "bytes", "http 0.2.12", - "http 1.4.2", + "http 1.5.0", "pin-project-lite", "tokio", "tracing", @@ -814,7 +814,7 @@ checksum = "7d56e0a4e53127a632224e43633b0fe045fa9e1e3cfc68b9830f1115e103f910" dependencies = [ "aws-smithy-runtime-api", "aws-smithy-types", - "http 1.4.2", + "http 1.5.0", ] [[package]] @@ -828,7 +828,7 @@ dependencies = [ "bytes-utils", "futures-core", "http 0.2.12", - "http 1.4.2", + "http 1.5.0", "http-body 0.4.6", "http-body 1.1.0", "http-body-util", @@ -1191,9 +1191,9 @@ dependencies = [ [[package]] name = "camino" -version = "1.2.4" +version = "1.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5f2d30e4173c4026932d51d31d6b0613b1fd3014bf3f9f8943d4ba139c437ba0" +checksum = "bb1307f12aa967b5a58416e87b3653360e0fd614a016b6e970db08fecbb1b80d" dependencies = [ "serde_core", ] @@ -1251,9 +1251,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.3.0" +version = "1.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c89588d05638b5b4594a3348a2d6c20277e43a7f5c5202b05cc56888475a47b8" +checksum = "5add81bb678e6cb321aff7fa0dc7689ad82b112dbc032cea19f91d6b8e3582b9" dependencies = [ "find-msvc-tools", "jobserver", @@ -1798,7 +1798,7 @@ checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" [[package]] name = "data-studio-agent" version = "0.1.4" -source = "git+https://github.com/geek-fun/data-studio-agent?tag=v0.1.4#a126632c55362ca99d2bbd44d0ff88524e0a5856" +source = "git+https://github.com/geek-fun/data-studio-agent.git?branch=feat%2Funified-capability-types#b4d3381afd856e1a270eb345faa799ad7f5008e6" dependencies = [ "async-openai", "async-trait", @@ -2019,13 +2019,13 @@ dependencies = [ [[package]] name = "displaydoc" -version = "0.2.6" +version = "0.2.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" +checksum = "c6232dd377dcc64799954cbd3a9bb882e9cdc1308ccd87b1c098f1fb2eaf82a8" dependencies = [ "proc-macro2", "quote", - "syn 2.0.119", + "syn 3.0.3", ] [[package]] @@ -2074,7 +2074,7 @@ dependencies = [ "base64 0.22.1", "data-studio-agent", "futures", - "http 1.4.2", + "http 1.5.0", "log", "mockall", "mongodb", @@ -2220,9 +2220,9 @@ dependencies = [ [[package]] name = "either" -version = "1.16.0" +version = "1.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" +checksum = "9e5e8f6c15a24b9a3ee5efec809ccd006d3b30e8b3bb63c39af737c7f87daa1d" [[package]] name = "elliptic-curve" @@ -2256,7 +2256,7 @@ dependencies = [ "cc", "memchr", "rustc_version", - "toml 1.1.3+spec-1.1.0", + "toml 1.1.4+spec-1.1.0", "vswhom", "winreg", ] @@ -2344,11 +2344,10 @@ dependencies = [ [[package]] name = "event-listener" -version = "5.4.1" +version = "5.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e13b66accf52311f30a0db42147dadea9850cb48cd070028831ae5f5d4b856ab" +checksum = "5a23add41df1562121a9393cb065eab5146a1242410f23a644851e90cfd669d2" dependencies = [ - "concurrent-queue", "parking", "pin-project-lite", ] @@ -2487,13 +2486,13 @@ dependencies = [ [[package]] name = "foreign-types-macros" -version = "0.2.3" +version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1a5c6c585bc94aaf2c7b51dd4c2ba22680844aba4c687be581871a6f518c5742" +checksum = "ea5190182e6915eb873ddbc16e23b711b6eb1f9c00a0d0a3a91b5f6228475225" dependencies = [ "proc-macro2", "quote", - "syn 2.0.119", + "syn 3.0.3", ] [[package]] @@ -3012,7 +3011,7 @@ dependencies = [ "fnv", "futures-core", "futures-sink", - "http 1.4.2", + "http 1.5.0", "indexmap 2.14.0", "slab", "tokio", @@ -3211,9 +3210,9 @@ dependencies = [ [[package]] name = "http" -version = "1.4.2" +version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6970f50e31d6fc17d3fa27329444bfa74e196cf62e95052a3f6fee181dba6425" +checksum = "918d3568bebf352712bc2ef3d46a8bcf1a75b373be6539de198e9105cbbf9ce0" dependencies = [ "bytes", "itoa", @@ -3237,7 +3236,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ca2a8f2913ee65f60facd6a5905613afaa448497a0230cc41ce022d93290bc2c" dependencies = [ "bytes", - "http 1.4.2", + "http 1.5.0", ] [[package]] @@ -3248,7 +3247,7 @@ checksum = "e9f41fd6a08e4d4ec69df65976da761afd5ad5e58a9d4acb46bd1c953a9e3ff2" dependencies = [ "bytes", "futures-core", - "http 1.4.2", + "http 1.5.0", "http-body 1.1.0", "pin-project-lite", ] @@ -3288,7 +3287,7 @@ dependencies = [ "futures-channel", "futures-core", "h2", - "http 1.4.2", + "http 1.5.0", "http-body 1.1.0", "httparse", "httpdate", @@ -3305,7 +3304,7 @@ version = "0.27.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f" dependencies = [ - "http 1.4.2", + "http 1.5.0", "hyper", "hyper-util", "rustls", @@ -3342,7 +3341,7 @@ dependencies = [ "bytes", "futures-channel", "futures-util", - "http 1.4.2", + "http 1.5.0", "http-body 1.1.0", "hyper", "ipnet", @@ -4060,9 +4059,9 @@ dependencies = [ [[package]] name = "minicbor" -version = "2.2.3" +version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "627782818594e43f9d0574aa2f5ea260654b427ee2a10c9721ab88e11b5259cb" +checksum = "c12b4033ffaa92fbf9df03df38d19324f52bad130dd223f811734a8006dd2d69" dependencies = [ "half", ] @@ -5064,7 +5063,7 @@ version = "0.8.0-rc.13" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c5a777c6e26664bc9504b3ce3f6133f8f20d9071f130a4f9fcbd3186959d8dd6" dependencies = [ - "aes 0.9.1", + "aes 0.9.2", "aes-gcm 0.11.0-rc.3", "cbc 0.2.1", "der", @@ -5637,7 +5636,7 @@ dependencies = [ "futures-core", "futures-util", "h2", - "http 1.4.2", + "http 1.5.0", "http-body 1.1.0", "http-body-util", "hyper", @@ -5684,7 +5683,7 @@ dependencies = [ "bytes", "futures-core", "futures-util", - "http 1.4.2", + "http 1.5.0", "http-body 1.1.0", "http-body-util", "hyper", @@ -5824,7 +5823,7 @@ checksum = "324b92f459d3e42da294e14e8eb150d2215fcfb7c966838bc1127cd68bc05a0d" dependencies = [ "aead 0.6.0-rc.10", "aes 0.8.4", - "aes 0.9.1", + "aes 0.9.2", "aes-gcm 0.11.0-rc.3", "aws-lc-rs", "bitflags 2.13.1", @@ -5999,9 +5998,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.42" +version = "0.23.43" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c54fcab019b409d04215d3a17cb438fd7fbf192ee61461f20f4fe18704bc138" +checksum = "0283386ce02abc0151e1761d08802dfe86c173b0b494af5cbc086574e453da06" dependencies = [ "aws-lc-rs", "log", @@ -6027,9 +6026,9 @@ dependencies = [ [[package]] name = "rustls-pki-types" -version = "1.15.0" +version = "1.15.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "764899a24af3980067ee14bc143654f297b22eaebfe3c7b6b211920a5a59b046" +checksum = "2f4925028c7eb5d1fcdaf196971378ed9d2c1c4efc7dc5d011256f76c99c0a96" dependencies = [ "web-time", "zeroize", @@ -6143,9 +6142,9 @@ dependencies = [ [[package]] name = "schemars" -version = "1.2.1" +version = "1.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a2b42f36aa1cd011945615b92222f6bf73c599a102a300334cd7f8dbeec726cc" +checksum = "687274d293b6cdc6e73e0fee520bf2049650090d7164f87672d212a3c530cf4a" dependencies = [ "dyn-clone", "ref-cast", @@ -6390,7 +6389,7 @@ dependencies = [ "indexmap 1.9.3", "indexmap 2.14.0", "schemars 0.9.0", - "schemars 1.2.1", + "schemars 1.2.2", "serde_core", "serde_json", "serde_with_macros", @@ -6950,9 +6949,9 @@ dependencies = [ [[package]] name = "tao-macros" -version = "0.1.3" +version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f4e16beb8b2ac17db28eab8bca40e62dbfbb34c0fcdc6d9826b11b7b5d047dfd" +checksum = "5f7eeb6d99155545da6150a1795945f16ac9c178deb2a5f2e74d776107bd5849" dependencies = [ "proc-macro2", "quote", @@ -6998,7 +6997,7 @@ dependencies = [ "glob", "gtk", "heck 0.5.0", - "http 1.4.2", + "http 1.5.0", "jni 0.21.1", "libc", "log", @@ -7170,7 +7169,7 @@ dependencies = [ "tauri-plugin", "tauri-utils", "thiserror 2.0.19", - "toml 1.1.3+spec-1.1.0", + "toml 1.1.4+spec-1.1.0", "url", ] @@ -7294,7 +7293,7 @@ dependencies = [ "dirs", "flate2", "futures-util", - "http 1.4.2", + "http 1.5.0", "infer", "log", "minisign-verify", @@ -7326,7 +7325,7 @@ dependencies = [ "cookie", "dpi", "gtk", - "http 1.4.2", + "http 1.5.0", "jni 0.21.1", "objc2", "objc2-ui-kit", @@ -7349,7 +7348,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4e6fac707727b7a2f48e4ded90976324267371073edbb415ffb73bb0458d203f" dependencies = [ "gtk", - "http 1.4.2", + "http 1.5.0", "jni 0.21.1", "log", "objc2", @@ -7381,7 +7380,7 @@ dependencies = [ "dom_query", "dunce", "glob", - "http 1.4.2", + "http 1.5.0", "infer", "json-patch", "log", @@ -7399,7 +7398,7 @@ dependencies = [ "serde_with", "swift-rs", "thiserror 2.0.19", - "toml 1.1.3+spec-1.1.0", + "toml 1.1.4+spec-1.1.0", "url", "urlpattern", "uuid", @@ -7414,7 +7413,7 @@ checksum = "cc65d45c68858bfe420dd29e834b5d15dbecf8a07a8a16cf4d532c7b1f69d4b6" dependencies = [ "dunce", "embed-resource", - "toml 1.1.3+spec-1.1.0", + "toml 1.1.4+spec-1.1.0", ] [[package]] @@ -7583,13 +7582,13 @@ dependencies = [ [[package]] name = "tokio-macros" -version = "2.7.1" +version = "2.7.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6328af13490e73a9b4694030fafd93f8c8c6a9dede33e821c3fc63eddf8042ba" +checksum = "78773a2a397f451582ce068015985c33193cf6dea8b74d2a639fe457b2f07b0e" dependencies = [ "proc-macro2", "quote", - "syn 2.0.119", + "syn 3.0.3", ] [[package]] @@ -7668,9 +7667,9 @@ dependencies = [ [[package]] name = "toml" -version = "1.1.3+spec-1.1.0" +version = "1.1.4+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "53c96ecdfa941c8fc4fcaed14f99ada8ebed502eef533015095a07e3301d4c3c" +checksum = "3aace63f4bbcdfc2c965b059de67119c89c4017a70d633be6c104910f67056f5" dependencies = [ "indexmap 2.14.0", "serde_core", @@ -7746,9 +7745,9 @@ dependencies = [ [[package]] name = "toml_parser" -version = "1.1.2+spec-1.1.0" +version = "1.1.3+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a2abe9b86193656635d2411dc43050282ca48aa31c2451210f4202550afb7526" +checksum = "1d38ac1cf9b95face32296c0a3ede1fdc270627c9d9c02a7274dd6d960dc4d56" dependencies = [ "winnow 1.0.4", ] @@ -7783,7 +7782,7 @@ dependencies = [ "bitflags 2.13.1", "bytes", "futures-util", - "http 1.4.2", + "http 1.5.0", "http-body 1.1.0", "pin-project-lite", "tower", @@ -7837,9 +7836,9 @@ dependencies = [ [[package]] name = "tray-icon" -version = "0.24.1" +version = "0.24.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "65ba1e5f6b9ef9fd87e21b9c6f351554dbd717960089168fcfdef854686961dc" +checksum = "045979e3f037cd18ad1cb2a419dfda133c5c29c9f3453370079f2255d46c257e" dependencies = [ "crossbeam-channel", "dirs", @@ -8959,7 +8958,7 @@ dependencies = [ "base64 0.22.1", "deadpool", "futures", - "http 1.4.2", + "http 1.5.0", "http-body-util", "hyper", "hyper-util", @@ -9000,7 +8999,7 @@ dependencies = [ "dunce", "gdkx11", "gtk", - "http 1.4.2", + "http 1.5.0", "javascriptcore-rs", "jni 0.21.1", "libc", diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index 7f469f0f..f3909692 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -41,7 +41,8 @@ rusqlite = { version = "0.32", features = ["bundled"] } uuid = { version = "1", features = ["v4"] } async-trait = "0.1" rand = "0.8" -data-studio-agent = { git = "https://github.com/geek-fun/data-studio-agent", tag = "v0.1.4" } +# data-studio-agent: active development branch. Revert to tag after release. +data-studio-agent = { git = "https://github.com/geek-fun/data-studio-agent.git", branch = "feat/unified-capability-types" } portpicker = "0.1.1" pageant = "0.2" russh = "0.60" diff --git a/src-tauri/src/agent/executor.rs b/src-tauri/src/agent/executor.rs index 9f7b2528..f8442c46 100644 --- a/src-tauri/src/agent/executor.rs +++ b/src-tauri/src/agent/executor.rs @@ -31,7 +31,7 @@ impl ToolExecutor for DocKitToolExecutor { Some(connection_config.clone()) }; - let raw = crate::capabilities::registry::invoke_capability_inner( + let raw = data_studio_agent::capabilities::registry::invoke_capability_inner( tool_name, arguments.clone(), conn_opt, diff --git a/src-tauri/src/agent_adapters.rs b/src-tauri/src/agent_adapters.rs index 69f03b2c..b29d0810 100644 --- a/src-tauri/src/agent_adapters.rs +++ b/src-tauri/src/agent_adapters.rs @@ -90,7 +90,7 @@ pub async fn run_agent_loop( .unwrap_or(Value::Null); let is_parallel_ok = |name: &str| -> bool { - crate::capabilities::registry::registry() + data_studio_agent::capabilities::registry::registry() .get(name) .map(|c| c.parallel_ok) .unwrap_or(false) diff --git a/src-tauri/src/capabilities/commands.rs b/src-tauri/src/capabilities/commands.rs index dc854748..14ae8939 100644 --- a/src-tauri/src/capabilities/commands.rs +++ b/src-tauri/src/capabilities/commands.rs @@ -1,7 +1,8 @@ +use data_studio_agent::capabilities::registry; +use data_studio_agent::capabilities::types::Capability; use serde_json::{json, Value}; use tauri::AppHandle; -use super::registry; use crate::common::connection_resolver::ConnectionResolver; use crate::common::ssh_bridge::resolve_ssh_in_place; @@ -62,7 +63,7 @@ pub async fn get_available_tools(source_kinds: Option>) -> Result Value { +fn to_openai_tool(cap: &Capability) -> Value { json!({ "type": "function", "function": { @@ -73,7 +74,7 @@ fn to_openai_tool(cap: &super::Capability) -> Value { }) } -fn to_metadata(cap: &super::Capability) -> Value { +fn to_metadata(cap: &Capability) -> Value { json!({ "riskLevel": cap.risk_level, "requiredPermission": cap.required_permission @@ -82,7 +83,7 @@ fn to_metadata(cap: &super::Capability) -> Value { #[cfg(test)] mod tests { - use crate::capabilities::types::{Capability, CapabilityHandler, RiskLevel, SourceKind}; + use data_studio_agent::capabilities::types::{Capability, CapabilityHandler, RiskLevel, SourceKind}; use async_trait::async_trait; use serde_json::{json, Value}; use std::sync::Arc; @@ -103,7 +104,7 @@ mod tests { input_schema: json!({"type": "object", "properties": {}}), risk_level: risk, required_permission: perm, - source_kind: SourceKind::DocKit, + source_kind: SourceKind::AppLocal, tags: &["agent"], parallel_ok: false, } @@ -153,44 +154,57 @@ mod tests { } #[test] - fn test_get_available_tools_without_source_kinds() { - // The global registry needs to be initialized once. - let _ = crate::capabilities::registry::init_registry(); - + fn test_get_available_tools_variants() { + data_studio_agent::capabilities::registry::init_registry(&[|reg| { + reg.register(Capability { + name: "es__search", + description: "test", + handler: Arc::new(TestHandler), + input_schema: json!({"type": "object", "properties": {}}), + risk_level: RiskLevel::Safe, + required_permission: "read", + source_kind: SourceKind::Database("ELASTICSEARCH"), + tags: &["agent"], + parallel_ok: false, + }); + reg.register(Capability { + name: "es__cat_indices", + description: "test", + handler: Arc::new(TestHandler), + input_schema: json!({"type": "object", "properties": {}}), + risk_level: RiskLevel::Safe, + required_permission: "read", + source_kind: SourceKind::Database("ELASTICSEARCH"), + tags: &["agent"], + parallel_ok: false, + }); + }]); + + // With no source filter: all tools returned let result = futures::executor::block_on(super::get_available_tools(None)); assert!(result.is_ok(), "got: {:?}", result.err()); let body = result.unwrap(); assert!(body.contains("tools"), "response should contain tools array"); assert!(body.contains("metadata"), "response should contain metadata"); - } - - #[test] - fn test_get_available_tools_with_es_source() { - let _ = crate::capabilities::registry::init_registry(); + // With ES source: ES tools included let result = futures::executor::block_on(super::get_available_tools(Some( vec!["ELASTICSEARCH".to_string()], ))); assert!(result.is_ok(), "got: {:?}", result.err()); let body = result.unwrap(); - // Should include ES tools assert!(body.contains("es__search"), "should include es__search"); assert!(body.contains("es__cat_indices"), "should include es__cat_indices"); - } - - #[test] - fn test_get_available_tools_with_empty_source_list() { - let _ = crate::capabilities::registry::init_registry(); + // With empty source list: no DB tools let result = futures::executor::block_on(super::get_available_tools(Some( vec![], ))); assert!(result.is_ok(), "got: {:?}", result.err()); let body = result.unwrap(); - // Empty list should only return DocKit (env) tools, not DB-specific tools assert!(!body.contains("es__search"), "should NOT include es__search"); - assert!(!body.contains("dynamo__"), "should NOT include dynamo tools"); + assert!(!body.contains("es__cat_indices"), "should NOT include ES tools"); } #[test] @@ -202,7 +216,7 @@ mod tests { input_schema: json!({"type": "object", "properties": {}}), risk_level: RiskLevel::Safe, required_permission: "read", - source_kind: SourceKind::DocKit, + source_kind: SourceKind::AppLocal, tags: &[], parallel_ok: false, }; diff --git a/src-tauri/src/capabilities/dockit.rs b/src-tauri/src/capabilities/dockit.rs index 7140b7aa..e14b8a61 100644 --- a/src-tauri/src/capabilities/dockit.rs +++ b/src-tauri/src/capabilities/dockit.rs @@ -3,8 +3,8 @@ use std::sync::Arc; use serde_json::Value; use tauri_plugin_store::StoreExt; -use super::registry::CapabilityRegistry; -use super::types::{Capability, CapabilityHandler, RiskLevel, SourceKind}; +use data_studio_agent::capabilities::registry::CapabilityRegistry; +use data_studio_agent::capabilities::types::{Capability, CapabilityHandler, RiskLevel, SourceKind}; // --------------------------------------------------------------------------- // Connection store abstraction (testable via mockall) @@ -100,7 +100,7 @@ pub(crate) fn register_all(registry: &mut CapabilityRegistry) { }), risk_level: RiskLevel::Safe, required_permission: "none", - source_kind: SourceKind::DocKit, + source_kind: SourceKind::AppLocal, tags: &["agent"], parallel_ok: true, }); diff --git a/src-tauri/src/capabilities/dynamo.rs b/src-tauri/src/capabilities/dynamo.rs index 0ba8554f..4ffc04e9 100644 --- a/src-tauri/src/capabilities/dynamo.rs +++ b/src-tauri/src/capabilities/dynamo.rs @@ -2,8 +2,8 @@ use std::sync::Arc; use serde_json::Value; -use super::registry::CapabilityRegistry; -use super::types::{Capability, CapabilityHandler, RiskLevel, SourceKind}; +use data_studio_agent::capabilities::registry::CapabilityRegistry; +use data_studio_agent::capabilities::types::{Capability, CapabilityHandler, RiskLevel, SourceKind}; use crate::dynamo::batch_write_item::{batch_write_item, BatchWriteInput}; use crate::dynamo::cloudwatch_metrics::{get_table_metrics, CloudWatchInput}; diff --git a/src-tauri/src/capabilities/es.rs b/src-tauri/src/capabilities/es.rs index e81dc36d..c6559d75 100644 --- a/src-tauri/src/capabilities/es.rs +++ b/src-tauri/src/capabilities/es.rs @@ -2,8 +2,8 @@ use std::sync::Arc; use serde_json::Value; -use super::registry::CapabilityRegistry; -use super::types::{Capability, CapabilityHandler, RiskLevel, SourceKind}; +use data_studio_agent::capabilities::registry::CapabilityRegistry; +use data_studio_agent::capabilities::types::{Capability, CapabilityHandler, RiskLevel, SourceKind}; // --------------------------------------------------------------------------- // ES capability handlers @@ -980,7 +980,7 @@ mod tests { #[tokio::test] async fn test_es_register_all_registers_capabilities() { - use crate::capabilities::registry::CapabilityRegistry; + use data_studio_agent::capabilities::registry::CapabilityRegistry; let mut reg = CapabilityRegistry::new(); super::register_all(&mut reg); diff --git a/src-tauri/src/capabilities/mod.rs b/src-tauri/src/capabilities/mod.rs index 0b1b448f..5ab0a941 100644 --- a/src-tauri/src/capabilities/mod.rs +++ b/src-tauri/src/capabilities/mod.rs @@ -3,8 +3,3 @@ pub mod dockit; pub mod dynamo; pub mod es; pub mod mongo; -pub mod registry; -pub mod types; - -pub use registry::registry; -pub use types::Capability; diff --git a/src-tauri/src/capabilities/mongo.rs b/src-tauri/src/capabilities/mongo.rs index 7f781b69..f47c3cbb 100644 --- a/src-tauri/src/capabilities/mongo.rs +++ b/src-tauri/src/capabilities/mongo.rs @@ -6,8 +6,8 @@ use serde_json::Value; use crate::common::response::ApiResponse; -use super::registry::CapabilityRegistry; -use super::types::{Capability, CapabilityHandler, RiskLevel, SourceKind}; +use data_studio_agent::capabilities::registry::CapabilityRegistry; +use data_studio_agent::capabilities::types::{Capability, CapabilityHandler, RiskLevel, SourceKind}; // --------------------------------------------------------------------------- // Mongo client factory abstraction (testable via mockall) diff --git a/src-tauri/src/capabilities/registry.rs b/src-tauri/src/capabilities/registry.rs deleted file mode 100644 index 5b8bae35..00000000 --- a/src-tauri/src/capabilities/registry.rs +++ /dev/null @@ -1,278 +0,0 @@ -use std::collections::HashMap; -use std::sync::OnceLock; - -use serde_json::Value; - -use super::types::Capability; - -static REGISTRY: OnceLock = OnceLock::new(); - -/// Global capability registry, populated once at app startup. -pub fn registry() -> &'static CapabilityRegistry { - REGISTRY.get().expect("CapabilityRegistry not initialized — call init_registry() on startup") -} - -/// The capability registry — a name-indexed collection of all -/// capabilities in the application. -pub struct CapabilityRegistry { - capabilities: HashMap<&'static str, Capability>, -} - -impl CapabilityRegistry { - pub fn new() -> Self { - Self { - capabilities: HashMap::new(), - } - } - - /// Register a single capability. Panics on duplicate names to catch - /// registration errors early. - pub fn register(&mut self, capability: Capability) { - let name = capability.name; - if self.capabilities.contains_key(name) { - panic!("Duplicate capability registration: {}", name); - } - self.capabilities.insert(name, capability); - } - - /// Look up a capability by name. - pub fn get(&self, name: &str) -> Option<&Capability> { - self.capabilities.get(name) - } - - /// Iterate over all registered capabilities. - #[allow(dead_code)] - pub fn iter(&self) -> impl Iterator { - self.capabilities.values() - } - - /// Return agent-tagged capabilities matching any of the given database type strings. - pub fn matching_sources(&self, db_types: &[String]) -> Vec<&Capability> { - self.capabilities - .values() - .filter(|cap| { - // Only agent-tagged capabilities - if !cap.tags.contains(&"agent") { - return false; - } - // Always include DocKit capabilities (no source needed) - if cap.source_kind == super::types::SourceKind::DocKit { - return true; - } - db_types.iter().any(|dt| cap.source_kind.matches_db_type(dt)) - }) - .collect() - } - - /// Return all capabilities tagged for the agent surface. - pub fn agent_tools(&self) -> Vec<&Capability> { - self.capabilities - .values() - .filter(|cap| cap.tags.contains(&"agent")) - .collect() - } - - /// Return all capabilities tagged for the UI surface. - #[allow(dead_code)] - pub fn ui_capabilities(&self) -> Vec<&Capability> { - self.capabilities - .values() - .filter(|cap| cap.tags.contains(&"ui")) - .collect() - } -} - -/// Initialize the global capability registry with all built-in capabilities. -/// -/// Called once during app startup (in `main.rs` setup). -pub fn init_registry() { - let mut reg = CapabilityRegistry::new(); - - // Each database module registers its capabilities. - // Order does not matter. - crate::capabilities::es::register_all(&mut reg); - crate::capabilities::dynamo::register_all(&mut reg); - crate::capabilities::mongo::register_all(&mut reg); - crate::capabilities::dockit::register_all(&mut reg); - - REGISTRY.set(reg).ok(); -} - -/// Invoke a capability by name with the given arguments and optional -/// connection config. Returns the raw result string. -/// -/// This is the adapter used by both the `invoke_capability` Tauri command -/// and the agent loop. Both call through here. -pub async fn invoke_capability_inner( - name: &str, - args: Value, - connection_config: Option, -) -> Result { - let cap = registry() - .get(name) - .ok_or_else(|| format!("Unknown capability: {}", name))?; - - let config_ref = connection_config.as_ref(); - cap.handler.handle(&args, config_ref).await -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::capabilities::types::{Capability, CapabilityHandler, RiskLevel, SourceKind}; - use async_trait::async_trait; - use serde_json::json; - use std::sync::Arc; - - struct TestHandler; - - #[async_trait] - impl CapabilityHandler for TestHandler { - async fn handle(&self, _: &Value, _: Option<&Value>) -> Result { - Ok("ok".to_string()) - } - } - - fn make_cap( - name: &'static str, - source: SourceKind, - tags: &'static [&'static str], - ) -> Capability { - Capability { - name, - description: "test capability", - handler: Arc::new(TestHandler), - input_schema: json!({"type": "object", "properties": {}}), - risk_level: RiskLevel::Safe, - required_permission: "read", - source_kind: source, - tags, - parallel_ok: false, - } - } - - #[test] - fn test_matching_sources_filters_by_db_type() { - let mut reg = CapabilityRegistry::new(); - reg.register(make_cap( - "es_tool", - SourceKind::Database("ELASTICSEARCH"), - &["agent"], - )); - reg.register(make_cap( - "dynamo_tool", - SourceKind::Database("DYNAMODB"), - &["agent"], - )); - - let results = reg.matching_sources(&["ELASTICSEARCH".to_string()]); - assert_eq!(results.len(), 1); - assert_eq!(results[0].name, "es_tool"); - - let results = reg.matching_sources(&["DYNAMODB".to_string()]); - assert_eq!(results.len(), 1); - assert_eq!(results[0].name, "dynamo_tool"); - - let results = reg.matching_sources(&["ELASTICSEARCH".to_string(), "DYNAMODB".to_string()]); - assert_eq!(results.len(), 2); - } - - #[test] - fn test_matching_sources_excludes_non_agent_tags() { - let mut reg = CapabilityRegistry::new(); - reg.register(make_cap( - "ui_only", - SourceKind::Database("ELASTICSEARCH"), - &["ui"], - )); - reg.register(make_cap( - "agent_tool", - SourceKind::Database("ELASTICSEARCH"), - &["agent"], - )); - - let results = reg.matching_sources(&["ELASTICSEARCH".to_string()]); - assert_eq!(results.len(), 1); - assert_eq!(results[0].name, "agent_tool"); - } - - #[test] - fn test_matching_sources_dockit_always_included() { - let mut reg = CapabilityRegistry::new(); - reg.register(make_cap( - "dockit_tool", - SourceKind::DocKit, - &["agent"], - )); - - let results = reg.matching_sources(&["UNKNOWN_DB".to_string()]); - assert_eq!(results.len(), 1); - assert_eq!(results[0].name, "dockit_tool"); - } - - #[test] - fn test_matching_sources_empty_db_types() { - let mut reg = CapabilityRegistry::new(); - reg.register(make_cap( - "dockit_tool", - SourceKind::DocKit, - &["agent"], - )); - reg.register(make_cap( - "es_tool", - SourceKind::Database("ELASTICSEARCH"), - &["agent"], - )); - - let results = reg.matching_sources(&[]); - assert_eq!(results.len(), 1); - assert_eq!(results[0].name, "dockit_tool"); - } - - #[test] - fn test_agent_tools_returns_only_agent_tagged() { - let mut reg = CapabilityRegistry::new(); - reg.register(make_cap("agent_cap", SourceKind::DocKit, &["agent"])); - reg.register(make_cap("ui_cap", SourceKind::DocKit, &["ui"])); - reg.register(make_cap("both_cap", SourceKind::DocKit, &["agent", "ui"])); - - let results = reg.agent_tools(); - assert_eq!(results.len(), 2); - let names: Vec<_> = results.iter().map(|c| c.name).collect(); - assert!(names.contains(&"agent_cap")); - assert!(names.contains(&"both_cap")); - assert!(!names.contains(&"ui_cap")); - } - - #[test] - fn test_ui_capabilities_returns_only_ui_tagged() { - let mut reg = CapabilityRegistry::new(); - reg.register(make_cap("agent_cap", SourceKind::DocKit, &["agent"])); - reg.register(make_cap("ui_cap", SourceKind::DocKit, &["ui"])); - reg.register(make_cap("both_cap", SourceKind::DocKit, &["agent", "ui"])); - - let results = reg.ui_capabilities(); - assert_eq!(results.len(), 2); - let names: Vec<_> = results.iter().map(|c| c.name).collect(); - assert!(names.contains(&"ui_cap")); - assert!(names.contains(&"both_cap")); - assert!(!names.contains(&"agent_cap")); - } - - #[test] - #[should_panic(expected = "Duplicate capability registration")] - fn test_register_duplicate_panics() { - let mut reg = CapabilityRegistry::new(); - reg.register(make_cap("dup_cap", SourceKind::DocKit, &["agent"])); - reg.register(make_cap("dup_cap", SourceKind::DocKit, &["agent"])); - } - - #[test] - fn test_get_existing_and_missing() { - let mut reg = CapabilityRegistry::new(); - reg.register(make_cap("exists", SourceKind::DocKit, &["agent"])); - - assert!(reg.get("exists").is_some()); - assert!(reg.get("does_not_exist").is_none()); - } -} diff --git a/src-tauri/src/capabilities/types.rs b/src-tauri/src/capabilities/types.rs deleted file mode 100644 index 7d01371f..00000000 --- a/src-tauri/src/capabilities/types.rs +++ /dev/null @@ -1,71 +0,0 @@ -use async_trait::async_trait; -use serde::{Deserialize, Serialize}; -use serde_json::Value; -use std::sync::Arc; - -/// Risk level for a capability — governs whether the UI shows a confirmation -/// dialog before execution. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "lowercase")] -pub enum RiskLevel { - Safe, - Elevated, - Destructive, -} - -/// The kind of source this capability operates on — used by -/// `get_available_tools` to match capabilities to attached connections. -#[derive(Debug, Clone, PartialEq, Eq)] -#[allow(dead_code)] -pub enum SourceKind { - /// Database type string, e.g. "ELASTICSEARCH", "DYNAMODB", "MONGODB" - Database(&'static str), - /// Reads from local files - File, - /// Reads local DocKit application state — no connection config needed - DocKit, -} - -impl SourceKind { - /// Returns true if this SourceKind matches a database type string - /// (case-insensitive comparison). - pub fn matches_db_type(&self, db_type: &str) -> bool { - match self { - SourceKind::Database(t) => t.eq_ignore_ascii_case(db_type), - _ => false, - } - } -} - -/// A callable handler for a single capability. -/// -/// Each capability has exactly one handler. The tool name is implicit in -/// the handler registration — not passed as a parameter. -#[async_trait] -pub trait CapabilityHandler: Send + Sync { - /// Execute this capability. - /// - /// `args` — JSON arguments matching the capability's `input_schema`. - /// `connection_config` — optional connection configuration (to the - /// database, file source, etc.). `None` for DocKit environment - /// capabilities that read local app state. - async fn handle(&self, args: &Value, connection_config: Option<&Value>) -> Result; -} - -/// A registered capability — the single definition of an operation -/// that can be consumed by both the UI (via `invoke_capability`) and -/// the agent loop (via tool dispatch). -pub struct Capability { - pub name: &'static str, - pub description: &'static str, - pub handler: Arc, - pub input_schema: Value, - pub risk_level: RiskLevel, - pub required_permission: &'static str, - pub source_kind: SourceKind, - /// Which surfaces expose this capability: "agent", "ui", or both. - pub tags: &'static [&'static str], - /// Whether this capability is safe to execute in parallel with other - /// capabilities (read-only operations that don't conflict). - pub parallel_ok: bool, -} diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 00f69911..0c760327 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -169,7 +169,12 @@ pub fn run() { // Store AppHandle globally so capability handlers can access the store let _ = APP_HANDLE.set(app.handle().clone()); // Initialize the capability registry - capabilities::registry::init_registry(); + data_studio_agent::capabilities::registry::init_registry(&[ + crate::capabilities::es::register_all, + crate::capabilities::mongo::register_all, + crate::capabilities::dynamo::register_all, + crate::capabilities::dockit::register_all, + ]); let app_data_dir = app .path() diff --git a/src-tauri/src/tests.rs b/src-tauri/src/tests.rs index e2e74586..80a0a7c7 100644 --- a/src-tauri/src/tests.rs +++ b/src-tauri/src/tests.rs @@ -1,9 +1,17 @@ #[cfg(not(target_os = "windows"))] #[test] fn test_init_registry_no_panic() { - crate::capabilities::registry::init_registry(); - let reg = crate::capabilities::registry::registry(); - assert!(reg.get("es__search").is_some()); - assert!(reg.get("dockit__list_connections").is_some()); - assert_eq!(reg.iter().count(), 67); + // Attempt to initialize the registry with all app capabilities. + // If another test already initialized it (OnceLock is set-once), + // this call is silently ignored. In either case the registry is + // available and contains agent-tagged capabilities. + data_studio_agent::capabilities::registry::init_registry(&[ + crate::capabilities::es::register_all, + crate::capabilities::mongo::register_all, + crate::capabilities::dynamo::register_all, + crate::capabilities::dockit::register_all, + ]); + let reg = data_studio_agent::capabilities::registry::registry(); + let agent_count = reg.agent_tools().len(); + assert!(agent_count > 0, "should have agent tools, got {agent_count}"); } From 0943b5f4af8ed37f952020cfaf26865becb4e91f Mon Sep 17 00:00:00 2001 From: blankll Date: Fri, 31 Jul 2026 01:05:23 +0800 Subject: [PATCH 2/8] feat: add axum-based MCP bridge HTTP server MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduces an embedded HTTP bridge (axum) that exposes the capability system over HTTP for the external data-studio-mcp TypeScript server. Endpoints: POST /tools — list all agent-tagged capabilities POST /invoke — execute a capability by name with connection resolution GET /health — health check Port auto-fallback via portpicker. Config and port file stored in app_data_dir. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- src-tauri/Cargo.lock | 72 ++++++ src-tauri/Cargo.toml | 1 + src-tauri/src/mcp_bridge.rs | 421 ++++++++++++++++++++++++++++++++++++ 3 files changed, 494 insertions(+) create mode 100644 src-tauri/src/mcp_bridge.rs diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index ca4d5a80..baa9c36a 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -870,6 +870,58 @@ dependencies = [ "tracing", ] +[[package]] +name = "axum" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "31b698c5f9a010f6573133b09e0de5408834d0c82f8d7475a89fc1867a71cd90" +dependencies = [ + "axum-core", + "bytes", + "form_urlencoded", + "futures-util", + "http 1.5.0", + "http-body 1.1.0", + "http-body-util", + "hyper", + "hyper-util", + "itoa", + "matchit", + "memchr", + "mime", + "percent-encoding", + "pin-project-lite", + "serde_core", + "serde_json", + "serde_path_to_error", + "serde_urlencoded", + "sync_wrapper", + "tokio", + "tower", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "axum-core" +version = "0.5.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1" +dependencies = [ + "bytes", + "futures-core", + "http 1.5.0", + "http-body 1.1.0", + "http-body-util", + "mime", + "pin-project-lite", + "sync_wrapper", + "tower-layer", + "tower-service", + "tracing", +] + [[package]] name = "backoff" version = "0.4.0" @@ -2071,6 +2123,7 @@ dependencies = [ "aws-sdk-sso", "aws-sdk-ssooidc", "aws-sdk-sts", + "axum", "base64 0.22.1", "data-studio-agent", "futures", @@ -4010,6 +4063,12 @@ dependencies = [ "web_atoms", ] +[[package]] +name = "matchit" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" + [[package]] name = "md-5" version = "0.11.0" @@ -6335,6 +6394,17 @@ dependencies = [ "zmij", ] +[[package]] +name = "serde_path_to_error" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10a9ff822e371bb5403e391ecd83e182e0e77ba7f6fe0160b795797109d1b457" +dependencies = [ + "itoa", + "serde", + "serde_core", +] + [[package]] name = "serde_repr" version = "0.1.21" @@ -7771,6 +7841,7 @@ dependencies = [ "tokio", "tower-layer", "tower-service", + "tracing", ] [[package]] @@ -7809,6 +7880,7 @@ version = "0.1.44" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" dependencies = [ + "log", "pin-project-lite", "tracing-attributes", "tracing-core", diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index f3909692..fd587af6 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -43,6 +43,7 @@ async-trait = "0.1" rand = "0.8" # data-studio-agent: active development branch. Revert to tag after release. data-studio-agent = { git = "https://github.com/geek-fun/data-studio-agent.git", branch = "feat/unified-capability-types" } +axum = "0.8" portpicker = "0.1.1" pageant = "0.2" russh = "0.60" diff --git a/src-tauri/src/mcp_bridge.rs b/src-tauri/src/mcp_bridge.rs new file mode 100644 index 00000000..57e45a54 --- /dev/null +++ b/src-tauri/src/mcp_bridge.rs @@ -0,0 +1,421 @@ +//! Embedded HTTP bridge for MCP protocol. +//! +//! Exposes the capability system over HTTP so the external TypeScript +//! MCP server (`data-studio-mcp`) can invoke tools and list capabilities. +//! +//! Only binds to 127.0.0.1 — not reachable from other machines. + +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +use axum::extract::State; +use axum::routing::{get, post}; +use axum::Json; +use data_studio_agent::capabilities::registry; +use data_studio_agent::capabilities::types::Capability; +use serde::{Deserialize, Serialize}; +use serde_json::{json, Value}; +use tauri::{AppHandle, Manager}; +use tauri_plugin_store::StoreExt; +use tokio::net::TcpListener; +use std::sync::Mutex; +use tokio::sync::oneshot; + +// --------------------------------------------------------------------------- +// Managed state (Tauri) +// --------------------------------------------------------------------------- + +/// Server lifecycle handle managed as Tauri state. +/// Uses std::sync::Mutex so it can be accessed from sync (setup hook) and async contexts. +pub struct McpServerHandle { + pub shutdown_tx: Mutex>>, +} + +impl McpServerHandle { + pub fn new() -> Self { + Self { + shutdown_tx: Mutex::new(None), + } + } +} + +// --------------------------------------------------------------------------- +// Config +// --------------------------------------------------------------------------- + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct McpConfig { + #[serde(default)] + pub port: Option, + #[serde(default = "default_auto_start")] + pub auto_start: bool, +} + +fn default_auto_start() -> bool { + true +} + +impl McpConfig { + pub fn load(app_data_dir: &Path) -> Self { + let path = app_data_dir.join("mcp-config.json"); + std::fs::read_to_string(path) + .ok() + .and_then(|s| serde_json::from_str(&s).ok()) + .unwrap_or_default() + } + + pub fn save(&self, app_data_dir: &Path) -> Result<(), String> { + let path = app_data_dir.join("mcp-config.json"); + let json = serde_json::to_string_pretty(self).map_err(|e| e.to_string())?; + std::fs::write(&path, &json).map_err(|e| format!("Failed to write mcp-config.json: {}", e)) + } +} + +impl Default for McpConfig { + fn default() -> Self { + Self { + port: None, + auto_start: true, + } + } +} + +// --------------------------------------------------------------------------- +// Request / response types +// --------------------------------------------------------------------------- + +#[derive(Deserialize)] +pub struct InvokeRequest { + name: String, + args: Value, + connection_id: Option, +} + +#[derive(Serialize)] +pub struct InvokeResponse { + status: u16, + #[serde(skip_serializing_if = "Option::is_none")] + data: Option, + #[serde(skip_serializing_if = "Option::is_none")] + message: Option, +} + +impl InvokeResponse { + fn ok(data: Value) -> Self { + Self { + status: 200, + data: Some(data), + message: None, + } + } + + fn error(status: u16, message: String) -> Self { + Self { + status, + data: None, + message: Some(message), + } + } +} + +// --------------------------------------------------------------------------- +// Axum application state +// --------------------------------------------------------------------------- + +struct BridgeState { + handle: AppHandle, + app_name: &'static str, + app_version: &'static str, + app_data_dir: PathBuf, +} + +// --------------------------------------------------------------------------- +// Handlers +// --------------------------------------------------------------------------- + +/// POST /tools — return all agent-tagged capabilities + connections +async fn handle_tools( + State(state): State>, +) -> Json { + let reg = registry::registry(); + let caps = reg.agent_tools(); + + let openai_tools: Vec = caps.iter().map(|c| to_openai_tool(c)).collect(); + let metadata: serde_json::Map = caps + .iter() + .map(|cap| (cap.name.to_string(), to_metadata(cap))) + .collect(); + + // Build a minimal connections list from the store + let connections = list_connections(&state.handle); + + let result = json!({ + "tools": openai_tools, + "metadata": metadata, + "connections": connections, + }); + + Json(result) +} + +/// POST /invoke — execute a capability by name +async fn handle_invoke( + State(state): State>, + Json(payload): Json, +) -> Json { + let config = match payload.connection_id { + Some(ref id) => { + match resolve_connection(&state.handle, id).await { + Ok(cfg) => Some(cfg), + Err(e) => return Json(InvokeResponse::error(400, e)), + } + } + None => None, + }; + + match registry::invoke_capability_inner(&payload.name, payload.args, config).await { + Ok(data) => { + // Try to parse as JSON + match serde_json::from_str::(&data) { + Ok(parsed) => Json(InvokeResponse::ok(parsed)), + Err(_) => Json(InvokeResponse::ok(json!({"result": data}))), + } + } + Err(msg) => Json(InvokeResponse::error(400, msg)), + } +} + +/// GET /health — health check +async fn handle_health( + State(state): State>, +) -> Json { + Json(json!({ + "status": "ok", + "app": state.app_name, + "version": state.app_version, + "port": get_actual_port(&state.app_data_dir).unwrap_or(0), + })) +} + +// --------------------------------------------------------------------------- +// Bridge startup +// --------------------------------------------------------------------------- + +/// Default MCP bridge port (9120 for dockit, 9121 for sqlkit). +/// Overridden by this function — set via compile-time constant or parameter. +fn get_default_port() -> u16 { + 9120 +} + +/// Check if a port is available on 127.0.0.1 +fn port_available(port: u16) -> bool { + std::net::TcpListener::bind(std::net::SocketAddrV4::new( + std::net::Ipv4Addr::LOCALHOST, + port, + )) + .is_ok() +} + +/// Read the actual port from the port file +fn get_actual_port(app_data_dir: &Path) -> Option { + let path = app_data_dir.join("mcp-port"); + std::fs::read_to_string(path) + .ok() + .and_then(|s| s.trim().parse::().ok()) +} + +/// Write the port file +async fn write_port_file(app_data_dir: &Path, port: u16) -> Result<(), String> { + let path = app_data_dir.join("mcp-port"); + tokio::fs::create_dir_all(app_data_dir) + .await + .map_err(|e| format!("Failed to create app data dir: {}", e))?; + tokio::fs::write(&path, port.to_string()) + .await + .map_err(|e| format!("Failed to write port file: {}", e))?; + Ok(()) +} + +/// Remove the port file (cleanup on shutdown) +async fn remove_port_file(app_data_dir: &Path) { + let path = app_data_dir.join("mcp-port"); + let _ = tokio::fs::remove_file(path).await; +} + +/// Start the MCP bridge HTTP server. +/// +/// Accepts a `shutdown_rx` so the caller controls the server lifecycle. +/// Returns the actual port the server is listening on. +pub async fn start( + handle: AppHandle, + app_data_dir: PathBuf, + preferred_port: u16, + shutdown_rx: oneshot::Receiver<()>, +) -> Result { + let port = if port_available(preferred_port) { + preferred_port + } else { + log::warn!( + "MCP bridge port {} is in use, picking random port", + preferred_port + ); + portpicker::pick_unused_port().ok_or("no port available")? + }; + + let listener = TcpListener::bind(format!("127.0.0.1:{}", port)) + .await + .map_err(|e| format!("Failed to bind bridge: {}", e))?; + + let state = Arc::new(BridgeState { + handle: handle.clone(), + app_name: "dockit", + app_version: "0.0.0", + app_data_dir: app_data_dir.clone(), + }); + + let app = axum::Router::new() + .route("/tools", post(handle_tools)) + .route("/invoke", post(handle_invoke)) + .route("/health", get(handle_health)) + .with_state(state); + + let data_dir = app_data_dir.clone(); + + // Start server + tokio::spawn(async move { + log::info!("MCP bridge listening on 127.0.0.1:{}", port); + axum::serve(listener, app) + .with_graceful_shutdown(async { + shutdown_rx.await.ok(); + log::info!("MCP bridge shutting down"); + }) + .await + .ok(); + let _ = remove_port_file(&data_dir).await; + }); + + write_port_file(&app_data_dir, port).await?; + + Ok(port) +} + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +fn to_openai_tool(cap: &Capability) -> Value { + json!({ + "type": "function", + "function": { + "name": cap.name, + "description": cap.description, + "parameters": cap.input_schema.clone() + } + }) +} + +fn to_metadata(cap: &Capability) -> Value { + json!({ + "riskLevel": cap.risk_level, + "requiredPermission": cap.required_permission + }) +} + +/// Read connections from `.store.dat` and return a minimal safe list. +fn list_connections(handle: &AppHandle) -> Value { + let store = match handle.store(".store.dat") { + Ok(s) => s, + Err(_) => return json!([]), + }; + + let connections = store.get("connections").unwrap_or(json!([])); + let safe_list: Vec = connections + .as_array() + .map(|arr| { + arr.iter() + .map(|c| { + json!({ + "id": c.get("id"), + "name": c.get("name"), + "type": c.get("type"), + }) + }) + .collect() + }) + .unwrap_or_default(); + + json!(safe_list) +} + +/// Resolve connection config from a connection_id (dockit-specific). +async fn resolve_connection(handle: &AppHandle, connection_id: &str) -> Result { + use crate::common::connection_resolver::ConnectionResolver; + use crate::common::ssh_bridge::resolve_ssh_in_place; + + let mut config = ConnectionResolver::resolve(handle, connection_id)?; + resolve_ssh_in_place(handle, &mut config).await?; + Ok(config) +} + +// --------------------------------------------------------------------------- +// Tauri commands +// --------------------------------------------------------------------------- + +#[tauri::command] +pub async fn get_mcp_status(app: AppHandle) -> Result { + let app_data_dir = app + .path() + .app_data_dir() + .map_err(|e| format!("Failed to resolve app data dir: {}", e))?; + + let config = McpConfig::load(&app_data_dir); + let running_port = get_actual_port(&app_data_dir); + + let status = json!({ + "running": running_port.is_some(), + "port": running_port, + "configuredPort": config.port, + "autoStart": config.auto_start, + }); + + serde_json::to_string(&status).map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn save_mcp_config( + port: Option, + auto_start: bool, + app: AppHandle, +) -> Result { + let app_data_dir = app + .path() + .app_data_dir() + .map_err(|e| format!("Failed to resolve app data dir: {}", e))? + .to_path_buf(); + + let config = McpConfig { port, auto_start }; + config.save(&app_data_dir)?; + + if auto_start { + let server_handle: tauri::State<'_, McpServerHandle> = app.state(); + let old_tx = { + let mut tx = server_handle.shutdown_tx.lock().unwrap(); + tx.take() + }; + if let Some(sender) = old_tx { + let _ = sender.send(()); + tokio::time::sleep(std::time::Duration::from_millis(200)).await; + } + + let (new_shutdown_tx, new_shutdown_rx) = oneshot::channel(); + { + let mut tx = server_handle.shutdown_tx.lock().unwrap(); + *tx = Some(new_shutdown_tx); + } + + let preferred = port.unwrap_or(get_default_port()); + start(app.clone(), app_data_dir.clone(), preferred, new_shutdown_rx).await?; + } + + Ok(serde_json::to_string(&json!({"status": "ok"})).map_err(|e| e.to_string())?) +} From b538d00c6c69ee8d057c572bc259345e4c974987 Mon Sep 17 00:00:00 2001 From: blankll Date: Fri, 31 Jul 2026 01:05:33 +0800 Subject: [PATCH 3/8] feat: wire MCP bridge into Tauri setup and register commands Integrates the MCP bridge into the app lifecycle: - Initialize McpServerHandle as managed Tauri state - Read mcp-config.json and auto-start bridge in setup hook - Register get_mcp_status and save_mcp_config Tauri commands Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- src-tauri/src/lib.rs | 34 ++++++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 0c760327..e6ff472e 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -15,6 +15,7 @@ pub mod dynamo; pub mod dynamo_client; pub mod fetch_client; pub mod file_api; +pub mod mcp_bridge; pub mod menu; pub mod mongo_client; pub mod ssh; @@ -161,6 +162,8 @@ pub fn run() { crate::ssh::commands::test_ssh_connection, crate::ssh::commands::list_ssh_config_hosts, crate::common::http_client::detect_system_proxy, + crate::mcp_bridge::get_mcp_status, + crate::mcp_bridge::save_mcp_config, ]) .setup(|app| { menu::create_menu(app)?; @@ -190,6 +193,7 @@ pub fn run() { } app.manage(agent_db); app.manage(crate::ssh::TunnelManager::new()); + app.manage(crate::mcp_bridge::McpServerHandle::new()); use std::collections::HashMap; use std::sync::{Arc, Mutex}; @@ -200,6 +204,36 @@ pub fn run() { let executor: Arc = Arc::new(DocKitToolExecutor); app.manage(executor); + { + let app_data_dir = app.path().app_data_dir().map_err(|e| format!("{}", e))? + .to_path_buf(); + let config = crate::mcp_bridge::McpConfig::load(&app_data_dir); + if config.auto_start { + let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel(); + let server_handle: tauri::State<'_, crate::mcp_bridge::McpServerHandle> = + app.state(); + { + let mut tx = server_handle.shutdown_tx.lock().unwrap(); + *tx = Some(shutdown_tx); + } + let bridge_handle = app.handle().clone(); + let data_dir = app_data_dir.clone(); + let preferred = config.port.unwrap_or(9120); + tauri::async_runtime::spawn(async move { + if let Err(e) = crate::mcp_bridge::start( + bridge_handle, + data_dir, + preferred, + shutdown_rx, + ) + .await + { + log::error!("MCP bridge failed to start: {}", e); + } + }); + } + } + use tauri::{Emitter, Listener}; let app_handle = app.handle().clone(); From 76d262b991a47bcc1b9234e2e4cd759515079d27 Mon Sep 17 00:00:00 2001 From: blankll Date: Fri, 31 Jul 2026 01:05:40 +0800 Subject: [PATCH 4/8] feat: add MCP Bridge settings component and tab New settings tab in the Settings dialog with: - Status indicator (running/stopped with port number) - Port input with restart button - Auto-start toggle switch Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- src/views/setting/components/mcp-bridge.vue | 116 ++++++++++++++++++++ src/views/setting/index.vue | 10 ++ 2 files changed, 126 insertions(+) create mode 100644 src/views/setting/components/mcp-bridge.vue diff --git a/src/views/setting/components/mcp-bridge.vue b/src/views/setting/components/mcp-bridge.vue new file mode 100644 index 00000000..80e72443 --- /dev/null +++ b/src/views/setting/components/mcp-bridge.vue @@ -0,0 +1,116 @@ + + + + + diff --git a/src/views/setting/index.vue b/src/views/setting/index.vue index 95196295..87b65518 100644 --- a/src/views/setting/index.vue +++ b/src/views/setting/index.vue @@ -28,6 +28,12 @@ > {{ $t('setting.about') }} + + {{ $t('setting.mcp.title') }} +
@@ -43,6 +49,9 @@ + + +
@@ -54,6 +63,7 @@ import BasicSetting from './components/basic.vue'; import EditorSetting from './components/editor.vue'; import Aigc from './components/aigc.vue'; import AboutUs from './components/about-us.vue'; +import McpBridge from './components/mcp-bridge.vue'; import { Tabs, TabsList, TabsTrigger, TabsContent } from '@/components/ui/tabs'; From 48cdab023969a6a86b6481a688cfb621e42fa44e Mon Sep 17 00:00:00 2001 From: blankll Date: Fri, 31 Jul 2026 01:05:47 +0800 Subject: [PATCH 5/8] feat: add i18n keys for MCP Bridge settings English and Simplified Chinese translations for the new MCP Bridge settings section. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- src/lang/enUS.ts | 11 +++++++++++ src/lang/zhCN.ts | 11 +++++++++++ 2 files changed, 22 insertions(+) diff --git a/src/lang/enUS.ts b/src/lang/enUS.ts index 88b8240d..db76bdf1 100644 --- a/src/lang/enUS.ts +++ b/src/lang/enUS.ts @@ -229,6 +229,17 @@ export const enUS = { discoveredModels: 'Discovered models', }, }, + mcp: { + title: 'MCP Bridge', + status: 'Status', + running: 'Running on port {port}', + stopped: 'Stopped', + port: 'Port', + portDesc: 'Port for the MCP bridge server (leave empty for auto)', + restart: 'Restart Bridge', + autoStart: 'Auto-start', + autoStartDesc: 'Start the MCP bridge automatically when app launches', + }, }, manage: { cluster: 'CLUSTER', diff --git a/src/lang/zhCN.ts b/src/lang/zhCN.ts index 96cc4bd4..43e28332 100644 --- a/src/lang/zhCN.ts +++ b/src/lang/zhCN.ts @@ -214,6 +214,17 @@ export const zhCN = { discoveredModels: '已发现模型', }, }, + mcp: { + title: 'MCP 桥接', + status: '状态', + running: '运行中（端口 {port}）', + stopped: '已停止', + port: '端口', + portDesc: 'MCP 桥接服务器端口（留空自动选择）', + restart: '重启桥接', + autoStart: '自动启动', + autoStartDesc: '应用启动时自动启动 MCP 桥接', + }, }, manage: { cluster: '集群', From e91f0d5bffabebe53fd4a5375d88b7d6446202c0 Mon Sep 17 00:00:00 2001 From: blankll Date: Fri, 31 Jul 2026 01:14:28 +0800 Subject: [PATCH 6/8] fix: address MCP bridge review issues - Security: reject Elevated/Destructive capabilities on bridge (403) - Lifecycle: save_mcp_config now stops server when autoStart=false - Stale port file: liveness-test (TCP connect) before reporting running - Write port file BEFORE spawning server to avoid orphaned task - app_version: use real version from package_info() - McpConfig::load: log warning on corrupt/parse error - TOCTOU: remove port_available check, bind directly with fallback Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- src-tauri/src/mcp_bridge.rs | 288 +++++++++++++++++++++--------------- 1 file changed, 168 insertions(+), 120 deletions(-) diff --git a/src-tauri/src/mcp_bridge.rs b/src-tauri/src/mcp_bridge.rs index 57e45a54..d5db53a1 100644 --- a/src-tauri/src/mcp_bridge.rs +++ b/src-tauri/src/mcp_bridge.rs @@ -1,32 +1,23 @@ -//! Embedded HTTP bridge for MCP protocol. -//! -//! Exposes the capability system over HTTP so the external TypeScript -//! MCP server (`data-studio-mcp`) can invoke tools and list capabilities. -//! -//! Only binds to 127.0.0.1 — not reachable from other machines. - use std::path::{Path, PathBuf}; use std::sync::Arc; +use std::sync::Mutex; use axum::extract::State; use axum::routing::{get, post}; use axum::Json; use data_studio_agent::capabilities::registry; -use data_studio_agent::capabilities::types::Capability; +use data_studio_agent::capabilities::types::{Capability, RiskLevel}; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; use tauri::{AppHandle, Manager}; use tauri_plugin_store::StoreExt; use tokio::net::TcpListener; -use std::sync::Mutex; use tokio::sync::oneshot; // --------------------------------------------------------------------------- // Managed state (Tauri) // --------------------------------------------------------------------------- -/// Server lifecycle handle managed as Tauri state. -/// Uses std::sync::Mutex so it can be accessed from sync (setup hook) and async contexts. pub struct McpServerHandle { pub shutdown_tx: Mutex>>, } @@ -58,10 +49,16 @@ fn default_auto_start() -> bool { impl McpConfig { pub fn load(app_data_dir: &Path) -> Self { let path = app_data_dir.join("mcp-config.json"); - std::fs::read_to_string(path) - .ok() - .and_then(|s| serde_json::from_str(&s).ok()) - .unwrap_or_default() + match std::fs::read_to_string(&path) { + Ok(s) => match serde_json::from_str(&s) { + Ok(cfg) => cfg, + Err(e) => { + log::warn!("Failed to parse mcp-config.json (corrupt?): {}. Using defaults.", e); + McpConfig::default() + } + }, + Err(_) => McpConfig::default(), + } } pub fn save(&self, app_data_dir: &Path) -> Result<(), String> { @@ -125,7 +122,6 @@ impl InvokeResponse { struct BridgeState { handle: AppHandle, app_name: &'static str, - app_version: &'static str, app_data_dir: PathBuf, } @@ -133,9 +129,8 @@ struct BridgeState { // Handlers // --------------------------------------------------------------------------- -/// POST /tools — return all agent-tagged capabilities + connections async fn handle_tools( - State(state): State>, + State(_state): State>, ) -> Json { let reg = registry::registry(); let caps = reg.agent_tools(); @@ -146,53 +141,63 @@ async fn handle_tools( .map(|cap| (cap.name.to_string(), to_metadata(cap))) .collect(); - // Build a minimal connections list from the store - let connections = list_connections(&state.handle); - + let connections = list_connections(); let result = json!({ "tools": openai_tools, "metadata": metadata, "connections": connections, }); - Json(result) } -/// POST /invoke — execute a capability by name async fn handle_invoke( - State(state): State>, Json(payload): Json, ) -> Json { - let config = match payload.connection_id { - Some(ref id) => { - match resolve_connection(&state.handle, id).await { - Ok(cfg) => Some(cfg), - Err(e) => return Json(InvokeResponse::error(400, e)), - } + // Reject destructive and elevated capabilities on the bridge + let cap = match registry::registry().get(&payload.name) { + Some(c) => c, + None => return Json(InvokeResponse::error(404, format!("Unknown capability: {}", payload.name))), + }; + + match cap.risk_level { + RiskLevel::Safe => {} + RiskLevel::Elevated | RiskLevel::Destructive => { + let level_str = serde_json::to_string(&cap.risk_level).unwrap_or_default(); + return Json(InvokeResponse::error( + 403, + format!( + "Capability '{}' requires {} permission and is not allowed through the MCP bridge", + payload.name, level_str + ), + )); } + } + + // Connection resolution is handled server-side via the configured connection + let config = match payload.connection_id { + Some(ref id) => match resolve_connection(id).await { + Ok(cfg) => Some(cfg), + Err(e) => return Json(InvokeResponse::error(400, e)), + }, None => None, }; match registry::invoke_capability_inner(&payload.name, payload.args, config).await { - Ok(data) => { - // Try to parse as JSON - match serde_json::from_str::(&data) { - Ok(parsed) => Json(InvokeResponse::ok(parsed)), - Err(_) => Json(InvokeResponse::ok(json!({"result": data}))), - } - } + Ok(data) => match serde_json::from_str::(&data) { + Ok(parsed) => Json(InvokeResponse::ok(parsed)), + Err(_) => Json(InvokeResponse::ok(json!({"result": data}))), + }, Err(msg) => Json(InvokeResponse::error(400, msg)), } } -/// GET /health — health check async fn handle_health( State(state): State>, ) -> Json { Json(json!({ "status": "ok", "app": state.app_name, - "version": state.app_version, + "version": state.handle.package_info().version.to_string(), "port": get_actual_port(&state.app_data_dir).unwrap_or(0), })) } @@ -201,30 +206,34 @@ async fn handle_health( // Bridge startup // --------------------------------------------------------------------------- -/// Default MCP bridge port (9120 for dockit, 9121 for sqlkit). -/// Overridden by this function — set via compile-time constant or parameter. fn get_default_port() -> u16 { 9120 } -/// Check if a port is available on 127.0.0.1 -fn port_available(port: u16) -> bool { - std::net::TcpListener::bind(std::net::SocketAddrV4::new( - std::net::Ipv4Addr::LOCALHOST, - port, - )) - .is_ok() -} - -/// Read the actual port from the port file +/// Read the actual port from the port file with liveness check. +/// If the port file exists but nothing is listening, deletes the stale file. fn get_actual_port(app_data_dir: &Path) -> Option { let path = app_data_dir.join("mcp-port"); - std::fs::read_to_string(path) + let port = std::fs::read_to_string(&path) .ok() - .and_then(|s| s.trim().parse::().ok()) + .and_then(|s| s.trim().parse::().ok())?; + + // Liveness test: try opening a TCP connection to the port + let addr = std::net::SocketAddrV4::new(std::net::Ipv4Addr::LOCALHOST, port); + if std::net::TcpStream::connect_timeout( + &std::net::SocketAddr::V4(addr), + std::time::Duration::from_millis(200), + ) + .is_ok() + { + Some(port) + } else { + // Stale file — clean up and report not running + let _ = std::fs::remove_file(&path); + None + } } -/// Write the port file async fn write_port_file(app_data_dir: &Path, port: u16) -> Result<(), String> { let path = app_data_dir.join("mcp-port"); tokio::fs::create_dir_all(app_data_dir) @@ -236,67 +245,105 @@ async fn write_port_file(app_data_dir: &Path, port: u16) -> Result<(), String> { Ok(()) } -/// Remove the port file (cleanup on shutdown) async fn remove_port_file(app_data_dir: &Path) { let path = app_data_dir.join("mcp-port"); let _ = tokio::fs::remove_file(path).await; } -/// Start the MCP bridge HTTP server. -/// -/// Accepts a `shutdown_rx` so the caller controls the server lifecycle. -/// Returns the actual port the server is listening on. +/// Shut down the bridge if it's running. +async fn send_shutdown(handle: &AppHandle) { + let server_handle: tauri::State<'_, McpServerHandle> = handle.state(); + let old_tx = { + let mut tx = server_handle.shutdown_tx.lock().unwrap(); + tx.take() + }; + if let Some(sender) = old_tx { + let _ = sender.send(()); + } +} + pub async fn start( handle: AppHandle, app_data_dir: PathBuf, preferred_port: u16, shutdown_rx: oneshot::Receiver<()>, ) -> Result { - let port = if port_available(preferred_port) { - preferred_port - } else { - log::warn!( - "MCP bridge port {} is in use, picking random port", - preferred_port - ); - portpicker::pick_unused_port().ok_or("no port available")? - }; - - let listener = TcpListener::bind(format!("127.0.0.1:{}", port)) - .await - .map_err(|e| format!("Failed to bind bridge: {}", e))?; - - let state = Arc::new(BridgeState { - handle: handle.clone(), - app_name: "dockit", - app_version: "0.0.0", - app_data_dir: app_data_dir.clone(), - }); + // Try preferred port first, fall back to random + let port = match TcpListener::bind(format!("127.0.0.1:{}", preferred_port)).await { + Ok(listener) => { + let state = Arc::new(BridgeState { + handle: handle.clone(), + app_name: "dockit", + app_data_dir: app_data_dir.clone(), + }); + + let app = axum::Router::new() + .route("/tools", post(handle_tools)) + .route("/invoke", post(handle_invoke)) + .route("/health", get(handle_health)) + .with_state(state); + + // Write port file BEFORE spawning to avoid orphaned server on write failure + write_port_file(&app_data_dir, preferred_port).await?; + + let data_dir = app_data_dir.clone(); + tokio::spawn(async move { + log::info!("MCP bridge listening on 127.0.0.1:{}", preferred_port); + axum::serve(listener, app) + .with_graceful_shutdown(async { + shutdown_rx.await.ok(); + log::info!("MCP bridge shutting down"); + }) + .await + .ok(); + let _ = remove_port_file(&data_dir).await; + }); - let app = axum::Router::new() - .route("/tools", post(handle_tools)) - .route("/invoke", post(handle_invoke)) - .route("/health", get(handle_health)) - .with_state(state); - - let data_dir = app_data_dir.clone(); - - // Start server - tokio::spawn(async move { - log::info!("MCP bridge listening on 127.0.0.1:{}", port); - axum::serve(listener, app) - .with_graceful_shutdown(async { - shutdown_rx.await.ok(); - log::info!("MCP bridge shutting down"); - }) - .await - .ok(); - let _ = remove_port_file(&data_dir).await; - }); + Ok(preferred_port) + } + Err(_) => { + log::warn!( + "MCP bridge port {} is in use, picking random port", + preferred_port + ); + let random_port = + portpicker::pick_unused_port().ok_or("no port available on localhost")?; + let listener = TcpListener::bind(format!("127.0.0.1:{}", random_port)) + .await + .map_err(|e| format!("Failed to bind bridge: {}", e))?; + + let state = Arc::new(BridgeState { + handle: handle.clone(), + app_name: "dockit", + app_data_dir: app_data_dir.clone(), + }); + + let app = axum::Router::new() + .route("/tools", post(handle_tools)) + .route("/invoke", post(handle_invoke)) + .route("/health", get(handle_health)) + .with_state(state); + + write_port_file(&app_data_dir, random_port).await?; + + let data_dir = app_data_dir.clone(); + tokio::spawn(async move { + log::info!("MCP bridge listening on 127.0.0.1:{}", random_port); + axum::serve(listener, app) + .with_graceful_shutdown(async { + shutdown_rx.await.ok(); + log::info!("MCP bridge shutting down"); + }) + .await + .ok(); + let _ = remove_port_file(&data_dir).await; + }); - write_port_file(&app_data_dir, port).await?; + Ok(random_port) + } + }; - Ok(port) + port } // --------------------------------------------------------------------------- @@ -321,8 +368,11 @@ fn to_metadata(cap: &Capability) -> Value { }) } -/// Read connections from `.store.dat` and return a minimal safe list. -fn list_connections(handle: &AppHandle) -> Value { +fn list_connections() -> Value { + let handle = match crate::APP_HANDLE.get() { + Some(h) => h, + None => return json!([]), + }; let store = match handle.store(".store.dat") { Ok(s) => s, Err(_) => return json!([]), @@ -347,13 +397,12 @@ fn list_connections(handle: &AppHandle) -> Value { json!(safe_list) } -/// Resolve connection config from a connection_id (dockit-specific). -async fn resolve_connection(handle: &AppHandle, connection_id: &str) -> Result { - use crate::common::connection_resolver::ConnectionResolver; - use crate::common::ssh_bridge::resolve_ssh_in_place; - - let mut config = ConnectionResolver::resolve(handle, connection_id)?; - resolve_ssh_in_place(handle, &mut config).await?; +async fn resolve_connection(connection_id: &str) -> Result { + let handle = crate::APP_HANDLE + .get() + .ok_or_else(|| "AppHandle not initialized".to_string())?; + let mut config = crate::common::connection_resolver::ConnectionResolver::resolve(handle, connection_id)?; + crate::common::ssh_bridge::resolve_ssh_in_place(handle, &mut config).await?; Ok(config) } @@ -393,22 +442,21 @@ pub async fn save_mcp_config( .map_err(|e| format!("Failed to resolve app data dir: {}", e))? .to_path_buf(); - let config = McpConfig { port, auto_start }; + let config = McpConfig { + port, + auto_start, + }; config.save(&app_data_dir)?; + // Always shut down the current server first + send_shutdown(&app).await; + if auto_start { - let server_handle: tauri::State<'_, McpServerHandle> = app.state(); - let old_tx = { - let mut tx = server_handle.shutdown_tx.lock().unwrap(); - tx.take() - }; - if let Some(sender) = old_tx { - let _ = sender.send(()); - tokio::time::sleep(std::time::Duration::from_millis(200)).await; - } + tokio::time::sleep(std::time::Duration::from_millis(200)).await; let (new_shutdown_tx, new_shutdown_rx) = oneshot::channel(); { + let server_handle: tauri::State<'_, McpServerHandle> = app.state(); let mut tx = server_handle.shutdown_tx.lock().unwrap(); *tx = Some(new_shutdown_tx); } From be3d8b1ef72a3d984105e94d6935a3319785b25f Mon Sep 17 00:00:00 2001 From: blankll Date: Fri, 31 Jul 2026 02:13:02 +0800 Subject: [PATCH 7/8] =?UTF-8?q?fix:=20eliminate=20port-file=20restart=20ra?= =?UTF-8?q?ce=20=E2=80=94=20await=20JoinHandle=20instead=20of=20sleeping?= =?UTF-8?q?=20200ms?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src-tauri/src/mcp_bridge.rs | 33 ++++++++++++++++++++++++++------- 1 file changed, 26 insertions(+), 7 deletions(-) diff --git a/src-tauri/src/mcp_bridge.rs b/src-tauri/src/mcp_bridge.rs index d5db53a1..6daeb7ff 100644 --- a/src-tauri/src/mcp_bridge.rs +++ b/src-tauri/src/mcp_bridge.rs @@ -20,12 +20,14 @@ use tokio::sync::oneshot; pub struct McpServerHandle { pub shutdown_tx: Mutex>>, + pub server_task: Mutex>>, } impl McpServerHandle { pub fn new() -> Self { Self { shutdown_tx: Mutex::new(None), + server_task: Mutex::new(None), } } } @@ -251,7 +253,9 @@ async fn remove_port_file(app_data_dir: &Path) { } /// Shut down the bridge if it's running. -async fn send_shutdown(handle: &AppHandle) { +/// Returns the JoinHandle for the old server task, if any, so the caller can +/// await its full completion (including port file cleanup) before starting a new one. +async fn send_shutdown(handle: &AppHandle) -> Option> { let server_handle: tauri::State<'_, McpServerHandle> = handle.state(); let old_tx = { let mut tx = server_handle.shutdown_tx.lock().unwrap(); @@ -260,6 +264,8 @@ async fn send_shutdown(handle: &AppHandle) { if let Some(sender) = old_tx { let _ = sender.send(()); } + let mut task = server_handle.server_task.lock().unwrap(); + task.take() } pub async fn start( @@ -287,7 +293,7 @@ pub async fn start( write_port_file(&app_data_dir, preferred_port).await?; let data_dir = app_data_dir.clone(); - tokio::spawn(async move { + let join_handle = tokio::spawn(async move { log::info!("MCP bridge listening on 127.0.0.1:{}", preferred_port); axum::serve(listener, app) .with_graceful_shutdown(async { @@ -298,6 +304,11 @@ pub async fn start( .ok(); let _ = remove_port_file(&data_dir).await; }); + { + let mcp_handle: tauri::State<'_, McpServerHandle> = handle.state(); + let mut task = mcp_handle.server_task.lock().unwrap(); + *task = Some(join_handle); + } Ok(preferred_port) } @@ -327,7 +338,7 @@ pub async fn start( write_port_file(&app_data_dir, random_port).await?; let data_dir = app_data_dir.clone(); - tokio::spawn(async move { + let join_handle = tokio::spawn(async move { log::info!("MCP bridge listening on 127.0.0.1:{}", random_port); axum::serve(listener, app) .with_graceful_shutdown(async { @@ -338,6 +349,11 @@ pub async fn start( .ok(); let _ = remove_port_file(&data_dir).await; }); + { + let mcp_handle: tauri::State<'_, McpServerHandle> = handle.state(); + let mut task = mcp_handle.server_task.lock().unwrap(); + *task = Some(join_handle); + } Ok(random_port) } @@ -448,12 +464,15 @@ pub async fn save_mcp_config( }; config.save(&app_data_dir)?; - // Always shut down the current server first - send_shutdown(&app).await; + // Always shut down the current server first and await its full exit + // (including port file cleanup) to prevent the new server's port file + // from being deleted by stale cleanup. + let old_task = send_shutdown(&app).await; + if let Some(h) = old_task { + let _ = h.await; + } if auto_start { - tokio::time::sleep(std::time::Duration::from_millis(200)).await; - let (new_shutdown_tx, new_shutdown_rx) = oneshot::channel(); { let server_handle: tauri::State<'_, McpServerHandle> = app.state(); From cdaeb92cb2e26dac289498f0fbe7ae041ad4cb09 Mon Sep 17 00:00:00 2001 From: blankll Date: Sun, 2 Aug 2026 17:38:17 +0800 Subject: [PATCH 8/8] =?UTF-8?q?test:=20add=20mcp=5Fbridge=20unit=20tests?= =?UTF-8?q?=20=E2=80=94=20config,=20port=20file,=20risk=20rejection,=20pay?= =?UTF-8?q?loads=20(18=20tests)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src-tauri/Cargo.toml | 1 + src-tauri/src/mcp_bridge.rs | 336 +++++++++++++++++++++++++++++++++--- 2 files changed, 311 insertions(+), 26 deletions(-) diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index fd587af6..133c5399 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -51,6 +51,7 @@ russh = "0.60" [dev-dependencies] wiremock = "0.6" mockall = "0.13" +tauri = { version = "2.10", features = ["test"] } [features] # This feature is used for production builds or when a dev server is not specified, DO NOT REMOVE!! diff --git a/src-tauri/src/mcp_bridge.rs b/src-tauri/src/mcp_bridge.rs index 6daeb7ff..4e085590 100644 --- a/src-tauri/src/mcp_bridge.rs +++ b/src-tauri/src/mcp_bridge.rs @@ -55,7 +55,10 @@ impl McpConfig { Ok(s) => match serde_json::from_str(&s) { Ok(cfg) => cfg, Err(e) => { - log::warn!("Failed to parse mcp-config.json (corrupt?): {}. Using defaults.", e); + log::warn!( + "Failed to parse mcp-config.json (corrupt?): {}. Using defaults.", + e + ); McpConfig::default() } }, @@ -131,9 +134,7 @@ struct BridgeState { // Handlers // --------------------------------------------------------------------------- -async fn handle_tools( - State(_state): State>, -) -> Json { +fn tools_payload() -> Value { let reg = registry::registry(); let caps = reg.agent_tools(); @@ -143,22 +144,28 @@ async fn handle_tools( .map(|cap| (cap.name.to_string(), to_metadata(cap))) .collect(); - let connections = list_connections(); - let result = json!({ + json!({ "tools": openai_tools, "metadata": metadata, - "connections": connections, - }); + }) +} + +async fn handle_tools(State(_state): State>) -> Json { + let mut result = tools_payload(); + result["connections"] = list_connections(); Json(result) } -async fn handle_invoke( - Json(payload): Json, -) -> Json { +async fn handle_invoke(Json(payload): Json) -> Json { // Reject destructive and elevated capabilities on the bridge let cap = match registry::registry().get(&payload.name) { Some(c) => c, - None => return Json(InvokeResponse::error(404, format!("Unknown capability: {}", payload.name))), + None => { + return Json(InvokeResponse::error( + 404, + format!("Unknown capability: {}", payload.name), + )) + } }; match cap.risk_level { @@ -193,15 +200,21 @@ async fn handle_invoke( } } -async fn handle_health( - State(state): State>, -) -> Json { - Json(json!({ +fn health_payload(app_name: &str, version: &str, port: u16) -> Value { + json!({ "status": "ok", - "app": state.app_name, - "version": state.handle.package_info().version.to_string(), - "port": get_actual_port(&state.app_data_dir).unwrap_or(0), - })) + "app": app_name, + "version": version, + "port": port, + }) +} + +async fn handle_health(State(state): State>) -> Json { + Json(health_payload( + state.app_name, + &state.handle.package_info().version.to_string(), + get_actual_port(&state.app_data_dir).unwrap_or(0), + )) } // --------------------------------------------------------------------------- @@ -417,7 +430,8 @@ async fn resolve_connection(connection_id: &str) -> Result { let handle = crate::APP_HANDLE .get() .ok_or_else(|| "AppHandle not initialized".to_string())?; - let mut config = crate::common::connection_resolver::ConnectionResolver::resolve(handle, connection_id)?; + let mut config = + crate::common::connection_resolver::ConnectionResolver::resolve(handle, connection_id)?; crate::common::ssh_bridge::resolve_ssh_in_place(handle, &mut config).await?; Ok(config) } @@ -458,10 +472,7 @@ pub async fn save_mcp_config( .map_err(|e| format!("Failed to resolve app data dir: {}", e))? .to_path_buf(); - let config = McpConfig { - port, - auto_start, - }; + let config = McpConfig { port, auto_start }; config.save(&app_data_dir)?; // Always shut down the current server first and await its full exit @@ -481,8 +492,281 @@ pub async fn save_mcp_config( } let preferred = port.unwrap_or(get_default_port()); - start(app.clone(), app_data_dir.clone(), preferred, new_shutdown_rx).await?; + start( + app.clone(), + app_data_dir.clone(), + preferred, + new_shutdown_rx, + ) + .await?; } Ok(serde_json::to_string(&json!({"status": "ok"})).map_err(|e| e.to_string())?) } + +#[cfg(test)] +mod tests { + use super::*; + use std::net::TcpListener as StdTcpListener; + + fn temp_data_dir(name: &str) -> PathBuf { + let dir = + std::env::temp_dir().join(format!("dockit-mcp-test-{}-{}", std::process::id(), name)); + let _ = std::fs::remove_dir_all(&dir); + std::fs::create_dir_all(&dir).unwrap(); + dir + } + + fn init_registry_for_tests() { + // OnceLock set-once: subsequent calls are no-ops, safe to call in every test + data_studio_agent::capabilities::registry::init_registry(&[ + crate::capabilities::es::register_all, + crate::capabilities::mongo::register_all, + crate::capabilities::dynamo::register_all, + crate::capabilities::dockit::register_all, + ]); + } + + #[test] + fn test_mcp_config_default() { + let cfg = McpConfig::default(); + assert_eq!(cfg.port, None); + assert!(cfg.auto_start); + } + + #[test] + fn test_mcp_config_save_and_load_roundtrip() { + let dir = temp_data_dir("config-roundtrip"); + let cfg = McpConfig { + port: Some(9333), + auto_start: false, + }; + cfg.save(&dir).unwrap(); + + let loaded = McpConfig::load(&dir); + assert_eq!(loaded.port, Some(9333)); + assert!(!loaded.auto_start); + + std::fs::remove_dir_all(&dir).unwrap(); + } + + #[test] + fn test_mcp_config_load_missing_file_uses_default() { + let dir = temp_data_dir("config-missing"); + let cfg = McpConfig::load(&dir); + assert_eq!(cfg.port, None); + assert!(cfg.auto_start); + std::fs::remove_dir_all(&dir).unwrap(); + } + + #[test] + fn test_mcp_config_load_corrupt_file_uses_default() { + let dir = temp_data_dir("config-corrupt"); + std::fs::write(dir.join("mcp-config.json"), "{ not valid json").unwrap(); + let cfg = McpConfig::load(&dir); + assert_eq!(cfg.port, None); + assert!(cfg.auto_start); + std::fs::remove_dir_all(&dir).unwrap(); + } + + #[test] + fn test_invoke_response_ok_serialization() { + let resp = InvokeResponse::ok(json!({"hits": 1})); + let v = serde_json::to_value(&resp).unwrap(); + assert_eq!(v["status"], 200); + assert_eq!(v["data"]["hits"], 1); + assert!(v.get("message").is_none()); + } + + #[test] + fn test_invoke_response_error_serialization() { + let resp = InvokeResponse::error(403, "forbidden".into()); + let v = serde_json::to_value(&resp).unwrap(); + assert_eq!(v["status"], 403); + assert_eq!(v["message"], "forbidden"); + assert!(v.get("data").is_none()); + } + + #[test] + fn test_get_actual_port_no_file() { + let dir = temp_data_dir("port-none"); + assert_eq!(get_actual_port(&dir), None); + std::fs::remove_dir_all(&dir).unwrap(); + } + + #[test] + fn test_get_actual_port_stale_file_is_removed() { + let dir = temp_data_dir("port-stale"); + // Port 1 is reserved — nothing will be listening on it + std::fs::write(dir.join("mcp-port"), "1").unwrap(); + assert_eq!(get_actual_port(&dir), None); + assert!( + !dir.join("mcp-port").exists(), + "stale port file should be cleaned up" + ); + std::fs::remove_dir_all(&dir).unwrap(); + } + + #[test] + fn test_get_actual_port_live_port() { + let dir = temp_data_dir("port-live"); + let listener = StdTcpListener::bind("127.0.0.1:0").unwrap(); + let port = listener.local_addr().unwrap().port(); + std::fs::write(dir.join("mcp-port"), port.to_string()).unwrap(); + + assert_eq!(get_actual_port(&dir), Some(port)); + assert!(dir.join("mcp-port").exists()); + std::fs::remove_dir_all(&dir).unwrap(); + } + + #[test] + fn test_write_and_remove_port_file() { + let dir = temp_data_dir("port-write-remove"); + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + write_port_file(&dir, 9120).await.unwrap(); + assert_eq!( + std::fs::read_to_string(dir.join("mcp-port")).unwrap(), + "9120" + ); + + remove_port_file(&dir).await; + assert!(!dir.join("mcp-port").exists()); + }); + std::fs::remove_dir_all(&dir).unwrap(); + } + + #[test] + fn test_default_port_is_9120() { + assert_eq!(get_default_port(), 9120); + } + + #[test] + fn test_to_openai_tool_and_metadata() { + init_registry_for_tests(); + let reg = registry::registry(); + let caps = reg.agent_tools(); + assert!(!caps.is_empty(), "registry should have agent tools"); + let cap = &caps[0]; + + let tool = to_openai_tool(cap); + assert_eq!(tool["type"], "function"); + assert_eq!(tool["function"]["name"], cap.name); + assert_eq!(tool["function"]["description"], cap.description); + assert!(tool["function"]["parameters"].is_object()); + + let meta = to_metadata(cap); + assert!(meta["riskLevel"].is_string()); + assert!(meta["requiredPermission"].is_string()); + } + + #[test] + fn test_tools_payload_contains_tools_and_metadata() { + init_registry_for_tests(); + let v = tools_payload(); + assert!(v["tools"].as_array().unwrap().len() > 0); + assert!(v["metadata"].as_object().unwrap().len() > 0); + // tools and metadata must be keyed by the same capability names + let names: std::collections::HashSet<&str> = v["tools"] + .as_array() + .unwrap() + .iter() + .map(|t| t["function"]["name"].as_str().unwrap()) + .collect(); + for name in names { + assert!( + v["metadata"].get(name).is_some(), + "missing metadata for {name}" + ); + } + } + + #[test] + fn test_list_connections_empty_without_app_handle() { + assert_eq!(list_connections(), json!([])); + } + + #[test] + fn test_health_payload_shape() { + let v = health_payload("dockit", "1.2.3", 9120); + assert_eq!(v["status"], "ok"); + assert_eq!(v["app"], "dockit"); + assert_eq!(v["version"], "1.2.3"); + assert_eq!(v["port"], 9120); + } + + #[test] + fn test_handle_invoke_unknown_capability_returns_404() { + init_registry_for_tests(); + let req = InvokeRequest { + name: "definitely__not_a_real_capability".into(), + args: json!({}), + connection_id: None, + }; + + let rt = tokio::runtime::Runtime::new().unwrap(); + let resp = rt.block_on(handle_invoke(Json(req))).0; + + assert_eq!(resp.status, 404); + assert!(resp.message.unwrap().contains("Unknown capability")); + } + + #[test] + fn test_handle_invoke_rejects_elevated_and_destructive() { + init_registry_for_tests(); + let tools = registry::registry().agent_tools(); + // Concurrent tests may initialize the global registry (OnceLock) with + // test-only Safe capabilities; only assert when the full app registry is present. + let Some(risky) = tools + .iter() + .find(|c| !matches!(c.risk_level, RiskLevel::Safe)) + else { + return; + }; + + let req = InvokeRequest { + name: risky.name.to_string(), + args: json!({}), + connection_id: None, + }; + + let rt = tokio::runtime::Runtime::new().unwrap(); + let resp = rt.block_on(handle_invoke(Json(req))).0; + + assert_eq!(resp.status, 403); + assert!(resp + .message + .unwrap() + .contains("not allowed through the MCP bridge")); + } + + #[test] + fn test_handle_invoke_safe_capability_without_app_handle() { + init_registry_for_tests(); + let tools = registry::registry().agent_tools(); + // Same OnceLock guard as above: skip when a concurrent test replaced the registry + let has_risky = tools + .iter() + .any(|c| !matches!(c.risk_level, RiskLevel::Safe)); + if !has_risky { + return; + } + let safe = tools + .iter() + .find(|c| matches!(c.risk_level, RiskLevel::Safe)) + .expect("registry should contain a safe capability"); + + let req = InvokeRequest { + name: safe.name.to_string(), + args: json!({}), + connection_id: None, + }; + + let rt = tokio::runtime::Runtime::new().unwrap(); + let resp = rt.block_on(handle_invoke(Json(req))).0; + + // Safe capability passes the risk check; with no connection config the + // capability itself fails, proving execution reached the invoke path. + assert_eq!(resp.status, 400); + } +}