diff --git a/Cargo.lock b/Cargo.lock index 66896e3..450f572 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -120,6 +120,12 @@ dependencies = [ "x11rb", ] +[[package]] +name = "ascii" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d92bec98840b8f03a5ff5413de5293bfcd8bf96467cf5452609f939ec6f5de16" + [[package]] name = "atomic-waker" version = "1.1.2" @@ -132,6 +138,18 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" +[[package]] +name = "base64" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8" + +[[package]] +name = "base64" +version = "0.21.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d297deb1925b89f2ccc13d7635fa0714f12c87adce1c75356b39ca9b7178567" + [[package]] name = "base64" version = "0.22.1" @@ -144,7 +162,7 @@ version = "0.72.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "993776b509cfb49c750f11b8f07a46fa23e0a1386ffc01fb1e7d343efc387895" dependencies = [ - "bitflags", + "bitflags 2.10.0", "cexpr", "clang-sys", "itertools 0.13.0", @@ -158,12 +176,27 @@ dependencies = [ "syn", ] +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + [[package]] name = "bitflags" version = "2.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "812e12b5285cc515a9c72a5c1d3b6d46a19dac5acfef5265968c166106e31dd3" +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + [[package]] name = "block2" version = "0.6.2" @@ -250,6 +283,12 @@ dependencies = [ "windows-link", ] +[[package]] +name = "chunked_transfer" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e4de3bc4ea267985becf712dc6d9eed8b04c953b3fcfb339ebc87acd9804901" + [[package]] name = "clang-sys" version = "1.8.1" @@ -350,6 +389,15 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + [[package]] name = "crc32fast" version = "1.5.0" @@ -380,7 +428,7 @@ version = "0.27.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f476fe445d41c9e991fd07515a6f463074b782242ccf4a5b7b1d1012e70824df" dependencies = [ - "bitflags", + "bitflags 2.10.0", "crossterm_winapi", "libc", "mio 0.8.11", @@ -405,6 +453,16 @@ version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + [[package]] name = "ctrlc" version = "3.5.1" @@ -416,6 +474,16 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + [[package]] name = "directories" version = "5.0.1" @@ -443,7 +511,7 @@ version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "89a09f22a6c6069a18470eb92d2298acf25463f14256d24778e1230d789a2aec" dependencies = [ - "bitflags", + "bitflags 2.10.0", "block2", "libc", "objc2", @@ -465,7 +533,7 @@ name = "drastis" version = "0.1.0" dependencies = [ "arboard", - "base64", + "base64 0.22.1", "chrono", "clap", "colored", @@ -478,26 +546,31 @@ dependencies = [ "llama-cpp-2", "log", "num_cpus", + "oauth2", + "open", "openai-agents-tracing", "openai_api_rust", "pulldown-cmark", "quick-xml", "r2d2", "r2d2_sqlite", + "rand 0.8.5", "ratatui", "ratatui-textarea", "regex", - "reqwest", + "reqwest 0.12.28", "roxmltree", "rusqlite", "rustyline", "serde", "serde_json", "serde_yaml", + "sha2", "similar", "tempfile", "termtree", "thiserror", + "tiny_http", "tokio", "tree-sitter", "tree-sitter-bash", @@ -515,6 +588,7 @@ dependencies = [ "tree-sitter-rust", "tree-sitter-toml-ng", "tree-sitter-typescript", + "url", "urlencoding", "walkdir", "zenpatch", @@ -775,6 +849,16 @@ dependencies = [ "slab", ] +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + [[package]] name = "gethostname" version = "1.1.0" @@ -801,8 +885,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "335ff9f135e4384c8150d6f27c6daed433577f86b4750418338c01a1a2528592" dependencies = [ "cfg-if", + "js-sys", "libc", "wasi", + "wasm-bindgen", ] [[package]] @@ -823,6 +909,25 @@ version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" +[[package]] +name = "h2" +version = "0.3.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0beca50380b1fc32983fc1cb4587bfa4bb9e78fc259aad4a0032d2080309222d" +dependencies = [ + "bytes", + "fnv", + "futures-core", + "futures-sink", + "futures-util", + "http 0.2.12", + "indexmap", + "slab", + "tokio", + "tokio-util", + "tracing", +] + [[package]] name = "h2" version = "0.4.13" @@ -834,7 +939,7 @@ dependencies = [ "fnv", "futures-core", "futures-sink", - "http", + "http 1.4.0", "indexmap", "slab", "tokio", @@ -915,6 +1020,17 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "http" +version = "0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "601cbb57e577e2f5ef5be8e7b83f0f63994f25aa94d673e54a92d5c516d101f1" +dependencies = [ + "bytes", + "fnv", + "itoa", +] + [[package]] name = "http" version = "1.4.0" @@ -925,6 +1041,17 @@ dependencies = [ "itoa", ] +[[package]] +name = "http-body" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ceab25649e9960c0311ea418d17bee82c0dcec1bd053b5f9a66e265a693bed2" +dependencies = [ + "bytes", + "http 0.2.12", + "pin-project-lite", +] + [[package]] name = "http-body" version = "1.0.1" @@ -932,7 +1059,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184" dependencies = [ "bytes", - "http", + "http 1.4.0", ] [[package]] @@ -943,8 +1070,8 @@ checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a" dependencies = [ "bytes", "futures-core", - "http", - "http-body", + "http 1.4.0", + "http-body 1.0.1", "pin-project-lite", ] @@ -954,6 +1081,36 @@ version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" +[[package]] +name = "httpdate" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" + +[[package]] +name = "hyper" +version = "0.14.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41dfc780fdec9373c01bae43289ea34c972e40ee3c9f6b3c8801a35f35586ce7" +dependencies = [ + "bytes", + "futures-channel", + "futures-core", + "futures-util", + "h2 0.3.27", + "http 0.2.12", + "http-body 0.4.6", + "httparse", + "httpdate", + "itoa", + "pin-project-lite", + "socket2 0.5.10", + "tokio", + "tower-service", + "tracing", + "want", +] + [[package]] name = "hyper" version = "1.8.1" @@ -964,9 +1121,9 @@ dependencies = [ "bytes", "futures-channel", "futures-core", - "h2", - "http", - "http-body", + "h2 0.4.13", + "http 1.4.0", + "http-body 1.0.1", "httparse", "itoa", "pin-project-lite", @@ -976,19 +1133,33 @@ dependencies = [ "want", ] +[[package]] +name = "hyper-rustls" +version = "0.24.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec3efd23720e2049821a693cbc7e65ea87c72f1c58ff2f9522ff332b1491e590" +dependencies = [ + "futures-util", + "http 0.2.12", + "hyper 0.14.32", + "rustls 0.21.12", + "tokio", + "tokio-rustls 0.24.1", +] + [[package]] name = "hyper-rustls" version = "0.27.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3c93eb611681b207e1fe55d5a71ecf91572ec8a6705cdb6857f7d8d5242cf58" dependencies = [ - "http", - "hyper", + "http 1.4.0", + "hyper 1.8.1", "hyper-util", - "rustls", + "rustls 0.23.36", "rustls-pki-types", "tokio", - "tokio-rustls", + "tokio-rustls 0.26.4", "tower-service", ] @@ -1000,7 +1171,7 @@ checksum = "70206fc6890eaca9fde8a0bf71caa2ddfc9fe045ac9e5c70df101a7dbde866e0" dependencies = [ "bytes", "http-body-util", - "hyper", + "hyper 1.8.1", "hyper-util", "native-tls", "tokio", @@ -1014,20 +1185,20 @@ version = "0.1.19" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "727805d60e7938b76b826a6ef209eb70eaa1812794f9424d4a4e2d740662df5f" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "futures-channel", "futures-core", "futures-util", - "http", - "http-body", - "hyper", + "http 1.4.0", + "http-body 1.0.1", + "hyper 1.8.1", "ipnet", "libc", "percent-encoding", "pin-project-lite", - "socket2", - "system-configuration", + "socket2 0.6.1", + "system-configuration 0.6.1", "tokio", "tower-service", "tracing", @@ -1211,6 +1382,25 @@ dependencies = [ "serde", ] +[[package]] +name = "is-docker" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "928bae27f42bc99b60d9ac7334e3a21d10ad8f1835a4e12ec3ec0464765ed1b3" +dependencies = [ + "once_cell", +] + +[[package]] +name = "is-wsl" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "173609498df190136aa7dea1a91db051746d339e18476eed5ca40521f02d7aa5" +dependencies = [ + "is-docker", + "once_cell", +] + [[package]] name = "is_terminal_polyfill" version = "1.70.2" @@ -1307,7 +1497,7 @@ version = "0.1.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d0b95e02c851351f877147b7deea7b1afb1df71b63aa5f8270716e0c5720616" dependencies = [ - "bitflags", + "bitflags 2.10.0", "libc", ] @@ -1476,7 +1666,7 @@ version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "71e2746dc3a24dd78b3cfcb7be93368c6de9963d30f43a6a73998a9cf4b17b46" dependencies = [ - "bitflags", + "bitflags 2.10.0", "cfg-if", "cfg_aliases", "libc", @@ -1488,7 +1678,7 @@ version = "0.30.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "74523f3a35e05aba87a1d978330aef40f67b0304ac79c1c00b294c9830543db6" dependencies = [ - "bitflags", + "bitflags 2.10.0", "cfg-if", "cfg_aliases", "libc", @@ -1523,6 +1713,26 @@ dependencies = [ "libc", ] +[[package]] +name = "oauth2" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c38841cdd844847e3e7c8d29cef9dcfed8877f8f56f9071f77843ecf3baf937f" +dependencies = [ + "base64 0.13.1", + "chrono", + "getrandom 0.2.16", + "http 0.2.12", + "rand 0.8.5", + "reqwest 0.11.27", + "serde", + "serde_json", + "serde_path_to_error", + "sha2", + "thiserror", + "url", +] + [[package]] name = "objc2" version = "0.6.3" @@ -1538,7 +1748,7 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d49e936b501e5c5bf01fda3a9452ff86dc3ea98ad5f283e1455153142d97518c" dependencies = [ - "bitflags", + "bitflags 2.10.0", "objc2", "objc2-core-graphics", "objc2-foundation", @@ -1550,7 +1760,7 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" dependencies = [ - "bitflags", + "bitflags 2.10.0", "dispatch2", "objc2", ] @@ -1561,7 +1771,7 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e022c9d066895efa1345f8e33e584b9f958da2fd4cd116792e15e07e4720a807" dependencies = [ - "bitflags", + "bitflags 2.10.0", "dispatch2", "objc2", "objc2-core-foundation", @@ -1580,7 +1790,7 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3e0adef53c21f888deb4fa59fc59f7eb17404926ee8a6f59f5df0fd7f9f3272" dependencies = [ - "bitflags", + "bitflags 2.10.0", "objc2", "objc2-core-foundation", ] @@ -1591,7 +1801,7 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "180788110936d59bab6bd83b6060ffdfffb3b922ba1396b312ae795e1de9d81d" dependencies = [ - "bitflags", + "bitflags 2.10.0", "objc2", "objc2-core-foundation", ] @@ -1608,13 +1818,24 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" +[[package]] +name = "open" +version = "5.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43bb73a7fa3799b198970490a51174027ba0d4ec504b03cd08caf513d40024bc" +dependencies = [ + "is-wsl", + "libc", + "pathdiff", +] + [[package]] name = "openai-agents-tracing" version = "0.1.0" dependencies = [ "anyhow", "chrono", - "reqwest", + "reqwest 0.12.28", "serde", "serde_json", "tokio", @@ -1641,7 +1862,7 @@ version = "0.10.75" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "08838db121398ad17ab8531ce9de97b244589089e290a384c900cb9ff7434328" dependencies = [ - "bitflags", + "bitflags 2.10.0", "cfg-if", "foreign-types", "libc", @@ -1714,6 +1935,12 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" +[[package]] +name = "pathdiff" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df94ce210e5bc13cb6651479fa48d14f601d9858cfe0467f43ae157023b938d3" + [[package]] name = "percent-encoding" version = "2.3.2" @@ -1744,7 +1971,7 @@ version = "0.18.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "97baced388464909d42d89643fe4361939af9b7ce7a31ee32a168f832a70f2a0" dependencies = [ - "bitflags", + "bitflags 2.10.0", "crc32fast", "fdeflate", "flate2", @@ -1809,7 +2036,7 @@ version = "0.10.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "76979bea66e7875e7509c4ec5300112b316af87fa7a252ca91c448b32dfe3993" dependencies = [ - "bitflags", + "bitflags 2.10.0", "getopts", "memchr", "pulldown-cmark-escape", @@ -1959,7 +2186,7 @@ version = "0.24.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0ebc917cfb527a566c37ecb94c7e3fd098353516fb4eb6bea17015ade0182425" dependencies = [ - "bitflags", + "bitflags 2.10.0", "cassowary", "crossterm", "indoc", @@ -1987,7 +2214,7 @@ version = "0.5.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" dependencies = [ - "bitflags", + "bitflags 2.10.0", ] [[package]] @@ -2030,24 +2257,65 @@ version = "0.8.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7a2d987857b319362043e95f5353c0535c1f58eec5336fdfcf626430af7def58" +[[package]] +name = "reqwest" +version = "0.11.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dd67538700a17451e7cba03ac727fb961abb7607553461627b97de0b89cf4a62" +dependencies = [ + "base64 0.21.7", + "bytes", + "encoding_rs", + "futures-core", + "futures-util", + "h2 0.3.27", + "http 0.2.12", + "http-body 0.4.6", + "hyper 0.14.32", + "hyper-rustls 0.24.2", + "ipnet", + "js-sys", + "log", + "mime", + "once_cell", + "percent-encoding", + "pin-project-lite", + "rustls 0.21.12", + "rustls-pemfile", + "serde", + "serde_json", + "serde_urlencoded", + "sync_wrapper 0.1.2", + "system-configuration 0.5.1", + "tokio", + "tokio-rustls 0.24.1", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", + "webpki-roots 0.25.4", + "winreg", +] + [[package]] name = "reqwest" version = "0.12.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "encoding_rs", "futures-channel", "futures-core", "futures-util", - "h2", - "http", - "http-body", + "h2 0.4.13", + "http 1.4.0", + "http-body 1.0.1", "http-body-util", - "hyper", - "hyper-rustls", + "hyper 1.8.1", + "hyper-rustls 0.27.7", "hyper-tls", "hyper-util", "js-sys", @@ -2060,7 +2328,7 @@ dependencies = [ "serde", "serde_json", "serde_urlencoded", - "sync_wrapper", + "sync_wrapper 1.0.2", "tokio", "tokio-native-tls", "tower", @@ -2098,7 +2366,7 @@ version = "0.32.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7753b721174eb8ff87a9a0e799e2d7bc3749323e773db92e0984debb00019d6e" dependencies = [ - "bitflags", + "bitflags 2.10.0", "fallible-iterator", "fallible-streaming-iterator", "hashlink", @@ -2118,13 +2386,25 @@ version = "1.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "146c9e247ccc180c1f61615433868c99f3de3ae256a30a43b49f67c2d9171f34" dependencies = [ - "bitflags", + "bitflags 2.10.0", "errno", "libc", "linux-raw-sys", "windows-sys 0.61.2", ] +[[package]] +name = "rustls" +version = "0.21.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f56a14d1f48b391359b22f731fd4bd7e43c97f3c50eee276f3aa09c94784d3e" +dependencies = [ + "log", + "ring", + "rustls-webpki 0.101.7", + "sct", +] + [[package]] name = "rustls" version = "0.23.36" @@ -2135,11 +2415,20 @@ dependencies = [ "once_cell", "ring", "rustls-pki-types", - "rustls-webpki", + "rustls-webpki 0.103.8", "subtle", "zeroize", ] +[[package]] +name = "rustls-pemfile" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c74cae0a4cf6ccbbf5f359f08efdf8ee7e1dc532573bf0db71968cb56b1448c" +dependencies = [ + "base64 0.21.7", +] + [[package]] name = "rustls-pki-types" version = "1.13.2" @@ -2149,6 +2438,16 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rustls-webpki" +version = "0.101.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b6275d1ee7a1cd780b64aca7726599a1dbc893b1e64144529e55c3c2f745765" +dependencies = [ + "ring", + "untrusted", +] + [[package]] name = "rustls-webpki" version = "0.103.8" @@ -2172,7 +2471,7 @@ version = "15.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2ee1e066dc922e513bda599c6ccb5f3bb2b0ea5870a579448f2622993f0a9a2f" dependencies = [ - "bitflags", + "bitflags 2.10.0", "cfg-if", "clipboard-win", "fd-lock", @@ -2227,13 +2526,23 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "sct" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da046153aa2352493d6cb7da4b6e5c0c057d8a1d0a9aa8560baffdd945acd414" +dependencies = [ + "ring", + "untrusted", +] + [[package]] name = "security-framework" version = "2.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "897b2245f0b511c87893af39b033e5ca9cce68824c4d7e7630b5a1d339658d02" dependencies = [ - "bitflags", + "bitflags 2.10.0", "core-foundation", "core-foundation-sys", "libc", @@ -2293,6 +2602,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_urlencoded" version = "0.7.1" @@ -2318,6 +2638,17 @@ dependencies = [ "unsafe-libyaml", ] +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + [[package]] name = "shlex" version = "1.3.0" @@ -2379,6 +2710,16 @@ version = "1.15.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" +[[package]] +name = "socket2" +version = "0.5.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e22376abed350d73dd1cd119b57ffccad95b4e585a7cda43e286245ce23c0678" +dependencies = [ + "libc", + "windows-sys 0.52.0", +] + [[package]] name = "socket2" version = "0.6.1" @@ -2446,6 +2787,12 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "sync_wrapper" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2047c6ded9c721764247e62cd3b03c09ffc529b2ba5b10ec482ae507a4a70160" + [[package]] name = "sync_wrapper" version = "1.0.2" @@ -2466,15 +2813,36 @@ dependencies = [ "syn", ] +[[package]] +name = "system-configuration" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba3a3adc5c275d719af8cb4272ea1c4a6d668a777f37e115f6d11ddbc1c8e0e7" +dependencies = [ + "bitflags 1.3.2", + "core-foundation", + "system-configuration-sys 0.5.0", +] + [[package]] name = "system-configuration" version = "0.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3c879d448e9d986b661742763247d3693ed13609438cf3d006f51f5368a5ba6b" dependencies = [ - "bitflags", + "bitflags 2.10.0", "core-foundation", - "system-configuration-sys", + "system-configuration-sys 0.6.0", +] + +[[package]] +name = "system-configuration-sys" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a75fb188eb626b924683e3b95e3a48e63551fcfb51949de2f06a9d91dbee93c9" +dependencies = [ + "core-foundation-sys", + "libc", ] [[package]] @@ -2540,6 +2908,18 @@ dependencies = [ "zune-jpeg 0.4.21", ] +[[package]] +name = "tiny_http" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "389915df6413a2e74fb181895f933386023c71110878cd0825588928e64cdc82" +dependencies = [ + "ascii", + "chunked_transfer", + "httpdate", + "log", +] + [[package]] name = "tinystr" version = "0.8.2" @@ -2562,7 +2942,7 @@ dependencies = [ "parking_lot", "pin-project-lite", "signal-hook-registry", - "socket2", + "socket2 0.6.1", "tokio-macros", "windows-sys 0.61.2", ] @@ -2588,13 +2968,23 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-rustls" +version = "0.24.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c28327cf380ac148141087fbfb9de9d7bd4e84ab5d2c28fbc911d753de8a7081" +dependencies = [ + "rustls 0.21.12", + "tokio", +] + [[package]] name = "tokio-rustls" version = "0.26.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61" dependencies = [ - "rustls", + "rustls 0.23.36", "tokio", ] @@ -2620,7 +3010,7 @@ dependencies = [ "futures-core", "futures-util", "pin-project-lite", - "sync_wrapper", + "sync_wrapper 1.0.2", "tokio", "tower-layer", "tower-service", @@ -2632,11 +3022,11 @@ version = "0.6.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d4e6559d53cc268e5031cd8429d05415bc4cb4aefc4aa5d6cc35fbf5b924a1f8" dependencies = [ - "bitflags", + "bitflags 2.10.0", "bytes", "futures-util", - "http", - "http-body", + "http 1.4.0", + "http-body 1.0.1", "iri-string", "pin-project-lite", "tower", @@ -2863,6 +3253,12 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" +[[package]] +name = "typenum" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "562d481066bde0658276a35467c4af00bdc6ee726305698a55b86e61d7ad82bb" + [[package]] name = "unicase" version = "2.9.0" @@ -2911,11 +3307,11 @@ version = "2.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "02d1a66277ed75f640d608235660df48c8e3c19f3b4edb6a263315626cc3c01d" dependencies = [ - "base64", + "base64 0.22.1", "flate2", "log", "once_cell", - "rustls", + "rustls 0.23.36", "rustls-pki-types", "serde", "serde_json", @@ -2933,6 +3329,7 @@ dependencies = [ "idna", "percent-encoding", "serde", + "serde_derive", ] [[package]] @@ -3086,6 +3483,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "webpki-roots" +version = "0.25.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f20c57d8d7db6d3b86154206ae5d8fba62dd39573114de97c2cb0578251f8e1" + [[package]] name = "webpki-roots" version = "0.26.11" @@ -3442,6 +3845,16 @@ version = "0.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" +[[package]] +name = "winreg" +version = "0.50.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "524e57b2c537c0f9b1e69f1965311ec12182b4122e45035b1508cd24d2adadb1" +dependencies = [ + "cfg-if", + "windows-sys 0.48.0", +] + [[package]] name = "wit-bindgen" version = "0.46.0" diff --git a/Cargo.toml b/Cargo.toml index 3d47522..483b5f6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -66,6 +66,12 @@ pulldown-cmark = "0.10" arboard = "3.6" base64 = "0.22" image = { version = "0.25", default-features = false, features = ["png", "jpeg"] } +oauth2 = "4.4" +tiny_http = "0.12" +open = "5.0" +sha2 = "0.10" +url = "2.5" +rand = "0.8" [dev-dependencies] diff --git a/cats/cat.py b/cats/cat.py index 2a70a83..091d89a 100644 --- a/cats/cat.py +++ b/cats/cat.py @@ -14,7 +14,7 @@ class Cat: def pet(self) -> str: if self.mood == "happy": return f"{self.name} purrs loudly." - return f"{self.name} tolerates the petting." + return f"{self.name} tolerates the petting with a calm blink." def begs(self) -> str: self.mood = "needy" diff --git a/src/domain/mod.rs b/src/domain/mod.rs index 3013724..18667f6 100644 --- a/src/domain/mod.rs +++ b/src/domain/mod.rs @@ -5,20 +5,20 @@ mod project; pub mod prompting; pub mod session; pub mod startup; +pub mod todo; pub mod tools; mod user_settings; pub mod workflow; -pub mod todo; pub use agent_mode::AgentModeType; pub use project::Project; pub use session::service::SessionService; pub use session::{Session, SessionRequest}; pub use startup::StartupService; -pub use user_settings::UserSettings; -pub use workflow::{CancellationToken, Chain}; #[allow(unused_imports)] -pub use todo::{TodoList, TodoItem}; +pub use todo::{TodoItem, TodoList}; +pub use user_settings::{AuthMethod, UserSettings}; +pub use workflow::{CancellationToken, Chain}; /// Model type enum matching the inference engine types #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -27,6 +27,13 @@ pub enum ModelType { OpenAI, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ModelAuthType { + Local, + ApiKey, + OAuth, +} + impl ModelType { pub fn as_str(&self) -> &'static str { match self { @@ -42,4 +49,29 @@ impl ModelType { _ => None, } } -} \ No newline at end of file +} + +impl ModelAuthType { + pub fn as_str(&self) -> &'static str { + match self { + ModelAuthType::Local => "local", + ModelAuthType::ApiKey => "api_key", + ModelAuthType::OAuth => "oauth", + } + } + + pub fn from_str(value: &str) -> Self { + match value { + "local" => ModelAuthType::Local, + "oauth" => ModelAuthType::OAuth, + _ => ModelAuthType::ApiKey, + } + } + + pub fn from_auth_method(method: &AuthMethod) -> Self { + match method { + AuthMethod::ApiKey => ModelAuthType::ApiKey, + AuthMethod::OAuth => ModelAuthType::OAuth, + } + } +} diff --git a/src/domain/permissions/checker.rs b/src/domain/permissions/checker.rs index 68d398b..7d3ce1d 100644 --- a/src/domain/permissions/checker.rs +++ b/src/domain/permissions/checker.rs @@ -1,8 +1,10 @@ use super::store::{PermissionStore, StoreError}; -use super::types::{PermissionConfig, PermissionDecision, PermissionRequest, PermissionScope}; +use super::types::{ + PermissionConfig, PermissionRequest, PermissionScope, SystemPermissionDecision, + UserPermissionDecision, +}; use crate::domain::session::Request; use crate::domain::tools::Tool; -use crate::utils::paths::is_within_root; use crate::utils::AskError; use std::path::PathBuf; use std::sync::Arc; @@ -23,7 +25,10 @@ pub struct PermissionChecker { } pub trait PermissionPrompter: Send + Sync { - fn ask_permission(&self, request: &PermissionRequest) -> Result; + fn ask_permission( + &self, + request: &PermissionRequest, + ) -> Result; } impl PermissionChecker { @@ -46,60 +51,68 @@ impl PermissionChecker { request: &dyn Request, project_id: Option, ) -> Result { + + // Auto-allow read-only operations within project root + if tool.is_read_only() { + let all_within_project = tool.get_affected_paths(request) + .iter() + .all(|p| p.starts_with(&request.project_root().to_path_buf())); + if all_within_project { + return Ok(true); + } + } + let permission_request = PermissionRequest::new( tool.name().to_string(), tool.get_command(request), tool.get_affected_paths(request), - super::types::PermissionScope::Project, + PermissionScope::Project, project_id, tool.is_read_only(), request.project_root().to_path_buf(), ); - let decision = self.resolve_permission(&permission_request)?; - if decision == PermissionDecision::Ask - && self.is_allowed_by_default(tool, request, &permission_request) - { - return Ok(true); - } - match decision { - PermissionDecision::AllowOnce => Ok(true), - PermissionDecision::Ask => self.prompt_and_store(&permission_request), - PermissionDecision::AllowAllReadsInSession - | PermissionDecision::AllowAllWritesInSession - | PermissionDecision::AllowCommandForProject => { - // These should have been resolved already, but handle them safely - Ok(true) + + // Dangerous commands always require a prompt + if let Some(ref cmd) = permission_request.command { + if self.is_dangerous_command(cmd) { + return self.prompt_and_store(&permission_request); } } + + // Restricted paths always require a prompt + if permission_request + .paths + .iter() + .any(|p| self.is_restricted_path(p)) + { + return self.prompt_and_store(&permission_request); + } + + let decision = self.resolve_permission(&permission_request)?; + match decision { + SystemPermissionDecision::Allow => Ok(true), + SystemPermissionDecision::Ask => self.prompt_and_store(&permission_request), + } } /// Store a permission decision fn store_permission_decision( &self, request: &PermissionRequest, - decision: PermissionDecision, + user_decision: UserPermissionDecision, ) -> Result<(), CheckerError> { - // Handle command-specific permissions (project-scoped, persistent) - if decision == PermissionDecision::AllowCommandForProject { - let permission = super::types::Permission::new( - request.tool_name.clone(), - request.command.clone(), // Store the specific command - None, // No resource pattern for commands - decision, - PermissionScope::Project, // Project scope = persistent - request.project_id, - ); - self.store.create_permission(permission)?; + if user_decision == UserPermissionDecision::AllowOnce { return Ok(()); } - // Handle session-wide permissions - if !matches!(decision, PermissionDecision::AllowAllReadsInSession | PermissionDecision::AllowAllWritesInSession) { - return Ok(()); - } - - // Get project root pattern for session-wide permissions - let resource_pattern = if let Some(first_path) = request.paths.first() { + // Minimal behavior: for project-scoped patch_files AlwaysAllow, store without + // resource pattern so it applies to any path within the project for this tool. + let resource_pattern = if request.tool_name == "patch_files" + && user_decision == UserPermissionDecision::AlwaysAllow + && request.scope == PermissionScope::Project + { + None + } else if let Some(first_path) = request.paths.first() { // Try to get the project root by going up the directory tree let mut path = first_path.clone(); while let Some(parent) = path.parent() { @@ -116,10 +129,10 @@ impl PermissionChecker { let permission = super::types::Permission::new( request.tool_name.clone(), - None, // No command pattern for session-wide + None, // No command pattern for session-wide resource_pattern, - decision, - PermissionScope::Session, + user_decision, + PermissionScope::Project, request.project_id, ); @@ -144,123 +157,38 @@ impl PermissionChecker { .any(|restricted| path_str.starts_with(restricted) || path_str.contains(restricted)) } - fn check_default_rules( - &self, - request: &PermissionRequest, - ) -> Result { - // Check for dangerous commands - if let Some(command) = &request.command { - if self.is_dangerous_command(command) { - return Ok(PermissionDecision::Ask); - } - } - - // Check for restricted paths - for path in &request.paths { - if self.is_restricted_path(path) { - return Ok(PermissionDecision::Ask); - } - } - - // Use default decision - Ok(self.config.default_decision.clone()) - } - - fn is_allowed_by_default( - &self, - tool: &dyn Tool, - request: &dyn Request, - permission_request: &PermissionRequest, - ) -> bool { - // Allow safe tools that don't require permission - if tool.skip_permission_check() { - return true; - } - - // Allow read-only tools if paths are within project root - if !tool.is_read_only() { - return false; - } - - let project_root = request.project_root(); - if permission_request.paths.is_empty() { - return false; - } - - permission_request.paths.iter().all(|path| { - let normalized = if path.is_absolute() { - path.to_path_buf() - } else { - project_root.join(path) - }; - is_within_root(&normalized, project_root) - }) - } - fn resolve_permission( &self, request: &PermissionRequest, - ) -> Result { + ) -> Result { let project_id: i32 = request.project_id.unwrap_or(0); if project_id == 0 { - return Ok(PermissionDecision::Ask); + return Ok(SystemPermissionDecision::Ask); } - // Check for command-specific permissions first (project-scoped, persistent) - if let Some(command) = &request.command { - if let Ok(Some(perm)) = self.store.find_permission( - &request.tool_name, - project_id, - command, - "", - ) { - if perm.decision == PermissionDecision::AllowCommandForProject - && perm.scope == PermissionScope::Project { - return Ok(PermissionDecision::AllowOnce); // Grant this specific operation - } - } + let command_str = request.command.as_ref().map_or("", |s| s.as_str()); + let path_str = request + .paths + .first() + .map(|p| p.to_string_lossy().to_string()) + .unwrap_or_default(); + let perm = + self.store + .find_permission(&request.tool_name, project_id, command_str, &path_str); + + if let Ok(Some(p)) = perm { + return Ok(p.system_decision()); } - // Check for session-wide permissions - let session_permissions = self.store.find_session_permissions(project_id)?; - for perm in session_permissions { - // CRITICAL SECURITY CHECK: Only grant session-wide permissions if ALL paths are within project root - let all_paths_in_project = request.paths.iter().all(|path| { - let normalized = if path.is_absolute() { - path.to_path_buf() - } else { - request.project_root.join(path) - }; - is_within_root(&normalized, &request.project_root) - }); - - if !all_paths_in_project { - // Paths outside project - don't use session permission, ask explicitly - continue; - } - - if request.is_read_only && perm.is_session_wide_all_reads() { - return Ok(PermissionDecision::AllowOnce); // Grant this specific operation - } - if !request.is_read_only && perm.is_session_wide_all_writes() { - return Ok(PermissionDecision::AllowOnce); // Grant this specific operation - } - } - - // Apply default security rules - self.check_default_rules(request) + Ok(self.config.default_decision.clone()) } fn prompt_and_store(&self, request: &PermissionRequest) -> Result { match self.prompter.ask_permission(request) { - Ok(decision @ PermissionDecision::AllowAllReadsInSession) - | Ok(decision @ PermissionDecision::AllowAllWritesInSession) - | Ok(decision @ PermissionDecision::AllowCommandForProject) => { + Ok(decision) => { self.store_permission_decision(request, decision)?; Ok(true) } - Ok(PermissionDecision::AllowOnce) => Ok(true), - Ok(PermissionDecision::Ask) => Ok(false), Err(AskError::IoError) => { Err(CheckerError::Failed("Permission prompt failed".to_string())) } @@ -272,7 +200,6 @@ impl PermissionChecker { mod tests { use super::*; use crate::domain::permissions::types::Permission; - use crate::domain::permissions::types::{PermissionConfig, PermissionDecision}; use crate::domain::permissions::{PermissionRequest, PermissionScope}; use crate::domain::session::{Request, SessionRequest}; use crate::domain::tools::{Error, Tool, ToolResult}; @@ -307,24 +234,6 @@ mod tests { Ok(permission) } - fn find_session_permissions(&self, project_id: i32) -> Result, StoreError> { - let permissions = self.permissions.lock().unwrap(); - let session_perms: Vec = permissions - .iter() - .filter(|p| { - p.project_id == Some(project_id) - && p.scope == PermissionScope::Session - && matches!( - p.decision, - PermissionDecision::AllowAllReadsInSession - | PermissionDecision::AllowAllWritesInSession - ) - }) - .cloned() - .collect(); - Ok(session_perms) - } - fn find_permission( &self, tool: &str, @@ -400,14 +309,14 @@ mod tests { struct TestPrompter { calls: Arc, - decision: PermissionDecision, + decision: UserPermissionDecision, } impl PermissionPrompter for TestPrompter { fn ask_permission( &self, _request: &PermissionRequest, - ) -> Result { + ) -> Result { self.calls.fetch_add(1, Ordering::SeqCst); Ok(self.decision.clone()) } @@ -598,12 +507,12 @@ mod tests { let calls = Arc::new(AtomicUsize::new(0)); let prompter = Arc::new(TestPrompter { calls: Arc::clone(&calls), - decision: PermissionDecision::AllowOnce, + decision: UserPermissionDecision::AlwaysAllow, }); let checker = PermissionChecker::new_with_prompter( store, PermissionConfig { - default_decision: PermissionDecision::Ask, + default_decision: SystemPermissionDecision::Ask, ..PermissionConfig::default() }, prompter, @@ -631,12 +540,12 @@ mod tests { let calls = Arc::new(AtomicUsize::new(0)); let prompter = Arc::new(TestPrompter { calls: Arc::clone(&calls), - decision: PermissionDecision::AllowOnce, + decision: UserPermissionDecision::AlwaysAllow, }); let checker = PermissionChecker::new_with_prompter( store, PermissionConfig { - default_decision: PermissionDecision::Ask, + default_decision: SystemPermissionDecision::Ask, ..PermissionConfig::default() }, prompter, @@ -663,12 +572,12 @@ mod tests { let calls = Arc::new(AtomicUsize::new(0)); let prompter = Arc::new(TestPrompter { calls: Arc::clone(&calls), - decision: PermissionDecision::AllowOnce, + decision: UserPermissionDecision::AlwaysAllow, }); let checker = PermissionChecker::new_with_prompter( store, PermissionConfig { - default_decision: PermissionDecision::Ask, + default_decision: SystemPermissionDecision::Ask, ..PermissionConfig::default() }, prompter, @@ -693,12 +602,12 @@ mod tests { let calls = Arc::new(AtomicUsize::new(0)); let prompter = Arc::new(TestPrompter { calls: Arc::clone(&calls), - decision: PermissionDecision::AllowOnce, + decision: UserPermissionDecision::AlwaysAllow, }); let checker = PermissionChecker::new_with_prompter( store, PermissionConfig { - default_decision: PermissionDecision::AllowOnce, + default_decision: SystemPermissionDecision::Allow, ..PermissionConfig::default() }, prompter, @@ -725,12 +634,12 @@ mod tests { let calls = Arc::new(AtomicUsize::new(0)); let prompter = Arc::new(TestPrompter { calls: Arc::clone(&calls), - decision: PermissionDecision::AllowOnce, + decision: UserPermissionDecision::AlwaysAllow, }); let checker = PermissionChecker::new_with_prompter( store, PermissionConfig { - default_decision: PermissionDecision::AllowOnce, + default_decision: SystemPermissionDecision::Allow, ..PermissionConfig::default() }, prompter, @@ -748,7 +657,7 @@ mod tests { } #[test] - fn resolve_permission_returns_ask_for_read_only() { + fn resolve_permission_returns_ask_when_no_stored_permission() { let abs = std::fs::canonicalize(".").expect("failed to canonicalize path"); let store = Arc::new(TestStore::with_permissions(vec![])); @@ -757,7 +666,7 @@ mod tests { PermissionConfig::default(), Arc::new(TestPrompter { calls: Arc::new(AtomicUsize::new(0)), - decision: PermissionDecision::AllowOnce, + decision: UserPermissionDecision::AlwaysAllow, }), ); let request = PermissionRequest::new( @@ -766,13 +675,13 @@ mod tests { vec![abs.clone()], PermissionScope::Project, Some(1), - true, // is_read_only + true, abs.parent().unwrap_or(&abs).to_path_buf(), ); let decision = checker.resolve_permission(&request).unwrap(); - assert_eq!(decision, PermissionDecision::Ask); + assert_eq!(decision, SystemPermissionDecision::Ask); } #[test] @@ -781,19 +690,21 @@ mod tests { let root = temp.path().to_path_buf(); let file_path = root.join("sample.txt"); + // Command permission with AllowOnce → system_decision = Ask let command_permission = Permission::new( "command_tool".to_string(), Some("rm -rf /".to_string()), None, - PermissionDecision::Ask, + UserPermissionDecision::AllowOnce, PermissionScope::Project, Some(1), ); + // Path permission with AlwaysAllow → system_decision = Allow let path_permission = Permission::new( "command_tool".to_string(), None, Some(file_path.to_string_lossy().to_string()), - PermissionDecision::AllowOnce, + UserPermissionDecision::AlwaysAllow, PermissionScope::Project, Some(1), ); @@ -806,7 +717,7 @@ mod tests { PermissionConfig::default(), Arc::new(TestPrompter { calls: Arc::new(AtomicUsize::new(0)), - decision: PermissionDecision::AllowOnce, + decision: UserPermissionDecision::AlwaysAllow, }), ); let request = PermissionRequest::new( @@ -815,13 +726,14 @@ mod tests { vec![file_path.clone()], PermissionScope::Project, Some(1), - false, // is_read_only + false, root, ); let decision = checker.resolve_permission(&request).unwrap(); - assert_eq!(decision, PermissionDecision::Ask); + // Command match has higher specificity, and AllowOnce → Ask + assert_eq!(decision, SystemPermissionDecision::Ask); } #[test] @@ -834,7 +746,7 @@ mod tests { "command_tool".to_string(), None, Some(file_path.to_string_lossy().to_string()), - PermissionDecision::AllowOnce, + UserPermissionDecision::AlwaysAllow, PermissionScope::Project, Some(1), ); @@ -844,7 +756,7 @@ mod tests { PermissionConfig::default(), Arc::new(TestPrompter { calls: Arc::new(AtomicUsize::new(0)), - decision: PermissionDecision::AllowOnce, + decision: UserPermissionDecision::AlwaysAllow, }), ); let request = PermissionRequest::new( @@ -853,13 +765,14 @@ mod tests { vec![file_path.clone()], PermissionScope::Project, Some(1), - false, // is_read_only + false, root, ); let decision = checker.resolve_permission(&request).unwrap(); - assert_eq!(decision, PermissionDecision::Ask); + // Path permission with AlwaysAllow → Allow + assert_eq!(decision, SystemPermissionDecision::Allow); } #[test] @@ -868,7 +781,7 @@ mod tests { "command_tool".to_string(), None, None, - PermissionDecision::Ask, + UserPermissionDecision::AllowOnce, PermissionScope::Project, Some(1), ); @@ -878,7 +791,7 @@ mod tests { PermissionConfig::default(), Arc::new(TestPrompter { calls: Arc::new(AtomicUsize::new(0)), - decision: PermissionDecision::AllowOnce, + decision: UserPermissionDecision::AlwaysAllow, }), ); let request = PermissionRequest::new( @@ -887,13 +800,14 @@ mod tests { vec![], PermissionScope::Project, Some(1), - false, // is_read_only + false, PathBuf::from("/tmp"), ); let decision = checker.resolve_permission(&request).unwrap(); - assert_eq!(decision, PermissionDecision::Ask); + // AllowOnce → system_decision = Ask + assert_eq!(decision, SystemPermissionDecision::Ask); } #[test] @@ -902,12 +816,12 @@ mod tests { let checker = PermissionChecker::new_with_prompter( store, PermissionConfig { - default_decision: PermissionDecision::AllowOnce, + default_decision: SystemPermissionDecision::Allow, ..PermissionConfig::default() }, Arc::new(TestPrompter { calls: Arc::new(AtomicUsize::new(0)), - decision: PermissionDecision::AllowOnce, + decision: UserPermissionDecision::AlwaysAllow, }), ); let request = PermissionRequest::new( @@ -916,147 +830,181 @@ mod tests { vec![], PermissionScope::Project, Some(1), - false, // is_read_only + false, PathBuf::from("/tmp"), ); let decision = checker.resolve_permission(&request).unwrap(); - assert_eq!(decision, PermissionDecision::AllowOnce); + assert_eq!(decision, SystemPermissionDecision::Allow); } #[test] - fn session_read_permission_only_works_within_project() { + fn stored_permission_does_not_apply_to_different_resource() { let project_dir = tempfile::tempdir().unwrap(); let external_dir = tempfile::tempdir().unwrap(); let project_root = project_dir.path().to_path_buf(); let external_file = external_dir.path().join("external.txt"); std::fs::write(&external_file, "data").unwrap(); - // Create a session-wide read permission - let session_perm = Permission::new( + // Create a permission with a specific resource pattern + let perm = Permission::new( "read_only".to_string(), None, Some(format!("{}/**", project_root.to_string_lossy())), - PermissionDecision::AllowAllReadsInSession, - PermissionScope::Session, + UserPermissionDecision::AlwaysAllow, + PermissionScope::Project, Some(1), ); - let store = Arc::new(TestStore::with_permissions(vec![session_perm])); + let store = Arc::new(TestStore::with_permissions(vec![perm])); let checker = PermissionChecker::new_with_prompter( store, PermissionConfig::default(), Arc::new(TestPrompter { calls: Arc::new(AtomicUsize::new(0)), - decision: PermissionDecision::AllowOnce, + decision: UserPermissionDecision::AlwaysAllow, }), ); - // Try to read a file OUTSIDE the project + // Request for a file that doesn't match the resource pattern let request = PermissionRequest::new( "read_only".to_string(), None, vec![external_file], PermissionScope::Project, Some(1), - true, // is_read_only + true, project_root, ); let decision = checker.resolve_permission(&request).unwrap(); - // Should NOT grant access via session permission - should Ask - assert_eq!(decision, PermissionDecision::Ask); + // Resource pattern doesn't match → falls back to default (Ask) + assert_eq!(decision, SystemPermissionDecision::Ask); } #[test] - fn session_write_permission_only_works_within_project() { + fn stored_permission_does_not_apply_to_different_project() { let project_dir = tempfile::tempdir().unwrap(); - let external_dir = tempfile::tempdir().unwrap(); let project_root = project_dir.path().to_path_buf(); - let external_file = external_dir.path().join("external.txt"); - std::fs::write(&external_file, "data").unwrap(); - // Create a session-wide write permission - let session_perm = Permission::new( + // Create a permission for project 1 + let perm = Permission::new( "write_tool".to_string(), None, - Some(format!("{}/**", project_root.to_string_lossy())), - PermissionDecision::AllowAllWritesInSession, - PermissionScope::Session, + None, + UserPermissionDecision::AlwaysAllow, + PermissionScope::Project, Some(1), ); - let store = Arc::new(TestStore::with_permissions(vec![session_perm])); + let store = Arc::new(TestStore::with_permissions(vec![perm])); let checker = PermissionChecker::new_with_prompter( store, PermissionConfig::default(), Arc::new(TestPrompter { calls: Arc::new(AtomicUsize::new(0)), - decision: PermissionDecision::AllowOnce, + decision: UserPermissionDecision::AlwaysAllow, }), ); - // Try to write a file OUTSIDE the project + // Request for project 2 - should not match let request = PermissionRequest::new( "write_tool".to_string(), None, - vec![external_file], + vec![project_root.join("file.txt")], PermissionScope::Project, - Some(1), - false, // is_read_only + Some(2), + false, project_root, ); let decision = checker.resolve_permission(&request).unwrap(); - // Should NOT grant access via session permission - should Ask - assert_eq!(decision, PermissionDecision::Ask); + assert_eq!(decision, SystemPermissionDecision::Ask); } #[test] - fn session_read_permission_works_within_project() { + fn stored_write_permission_grants_access() { + let project_dir = tempfile::tempdir().unwrap(); + let project_root = project_dir.path().to_path_buf(); + + // Create an AlwaysAllow permission for write_tool (no resource pattern = matches any) + let perm = Permission::new( + "write_tool".to_string(), + None, + None, + UserPermissionDecision::AlwaysAllow, + PermissionScope::Project, + Some(1), + ); + + let store = Arc::new(TestStore::with_permissions(vec![perm])); + let checker = PermissionChecker::new_with_prompter( + store, + PermissionConfig::default(), + Arc::new(TestPrompter { + calls: Arc::new(AtomicUsize::new(0)), + decision: UserPermissionDecision::AlwaysAllow, + }), + ); + + let request = PermissionRequest::new( + "write_tool".to_string(), + None, + vec![project_root.join("internal.txt")], + PermissionScope::Project, + Some(1), + false, + project_root, + ); + + let decision = checker.resolve_permission(&request).unwrap(); + + assert_eq!(decision, SystemPermissionDecision::Allow); + } + + #[test] + fn stored_read_permission_grants_access() { let project_dir = tempfile::tempdir().unwrap(); let project_root = project_dir.path().to_path_buf(); let internal_file = project_root.join("internal.txt"); std::fs::write(&internal_file, "data").unwrap(); - // Create a session-wide read permission - let session_perm = Permission::new( + // Create an AlwaysAllow permission for read_only (no resource pattern = matches any) + let perm = Permission::new( "read_only".to_string(), None, - Some(format!("{}/**", project_root.to_string_lossy())), - PermissionDecision::AllowAllReadsInSession, - PermissionScope::Session, + None, + UserPermissionDecision::AlwaysAllow, + PermissionScope::Project, Some(1), ); - let store = Arc::new(TestStore::with_permissions(vec![session_perm])); + let store = Arc::new(TestStore::with_permissions(vec![perm])); let checker = PermissionChecker::new_with_prompter( store, PermissionConfig::default(), Arc::new(TestPrompter { calls: Arc::new(AtomicUsize::new(0)), - decision: PermissionDecision::AllowOnce, + decision: UserPermissionDecision::AlwaysAllow, }), ); - // Try to read a file INSIDE the project let request = PermissionRequest::new( "read_only".to_string(), None, vec![internal_file], PermissionScope::Project, Some(1), - true, // is_read_only + true, project_root, ); let decision = checker.resolve_permission(&request).unwrap(); - // SHOULD grant access via session permission - assert_eq!(decision, PermissionDecision::AllowOnce); + assert_eq!(decision, SystemPermissionDecision::Allow); } #[test] @@ -1064,12 +1012,12 @@ mod tests { let project_dir = tempfile::tempdir().unwrap(); let project_root = project_dir.path().to_path_buf(); - // Create a project-scoped command permission + // Create a project-scoped command permission with AlwaysAllow let command_perm = Permission::new( "shell_exec".to_string(), Some("npm test".to_string()), None, - PermissionDecision::AllowCommandForProject, + UserPermissionDecision::AlwaysAllow, PermissionScope::Project, Some(1), ); @@ -1080,7 +1028,7 @@ mod tests { PermissionConfig::default(), Arc::new(TestPrompter { calls: Arc::new(AtomicUsize::new(0)), - decision: PermissionDecision::AllowOnce, + decision: UserPermissionDecision::AlwaysAllow, }), ); @@ -1091,12 +1039,12 @@ mod tests { vec![], PermissionScope::Project, Some(1), - false, // is_read_only + false, project_root.clone(), ); let decision = checker.resolve_permission(&request).unwrap(); - assert_eq!(decision, PermissionDecision::AllowOnce); // Allowed via stored permission + assert_eq!(decision, SystemPermissionDecision::Allow); // Request for a DIFFERENT command should ask let different_request = PermissionRequest::new( @@ -1110,6 +1058,48 @@ mod tests { ); let decision = checker.resolve_permission(&different_request).unwrap(); - assert_eq!(decision, PermissionDecision::Ask); // Different command, must ask + assert_eq!(decision, SystemPermissionDecision::Ask); + } + + #[test] + fn patch_files_permission_within_project_root_applies_to_other_project_files() { + let project_dir = tempfile::tempdir().unwrap(); + let project_root = project_dir.path().to_path_buf(); + + // User previously created an AlwaysAllow permission for patch_files + // targeting one file inside this project. + let existing_perm = Permission::new( + "patch_files".to_string(), + None, + None, + UserPermissionDecision::AlwaysAllow, + PermissionScope::Project, + Some(1), + ); + + let store = Arc::new(TestStore::with_permissions(vec![existing_perm])); + let checker = PermissionChecker::new_with_prompter( + store, + PermissionConfig::default(), + Arc::new(TestPrompter { + calls: Arc::new(AtomicUsize::new(0)), + decision: UserPermissionDecision::AlwaysAllow, + }), + ); + + // New request patches a DIFFERENT file, but still within the same project root. + // Expected behavior (desired): this should be allowed by the existing permission. + let request = PermissionRequest::new( + "patch_files".to_string(), + None, + vec![project_root.join("src/b.rs")], + PermissionScope::Project, + Some(1), + false, + project_root, + ); + + let decision = checker.resolve_permission(&request).unwrap(); + assert_eq!(decision, SystemPermissionDecision::Allow); } } \ No newline at end of file diff --git a/src/domain/permissions/store.rs b/src/domain/permissions/store.rs index 3e57416..dd94ade 100644 --- a/src/domain/permissions/store.rs +++ b/src/domain/permissions/store.rs @@ -1,4 +1,4 @@ -use super::types::{Permission, PermissionDecision, PermissionScope}; +use super::types::{Permission, PermissionScope, UserPermissionDecision}; use crate::infrastructure::db::DbPool; use rusqlite::params; use thiserror::Error; @@ -18,7 +18,6 @@ pub trait PermissionStore { command_pattern: &str, resource_pattern: &str, ) -> Result, StoreError>; - fn find_session_permissions(&self, project_id: i32) -> Result, StoreError>; } pub struct SqlitePermissionStore { @@ -31,18 +30,6 @@ impl SqlitePermissionStore { } fn row_to_permission(&self, row: &rusqlite::Row) -> Result { - let decision_str: String = row.get("decision")?; - let decision = match decision_str.as_str() { - "allow_all_reads_session" => PermissionDecision::AllowAllReadsInSession, - "allow_all_writes_session" => PermissionDecision::AllowAllWritesInSession, - "allow_command_for_project" => PermissionDecision::AllowCommandForProject, - "ask" => PermissionDecision::Ask, - "once" => PermissionDecision::AllowOnce, - // Handle legacy values gracefully - "allow" | "deny" => PermissionDecision::Ask, - _ => PermissionDecision::Ask, // Default to ask on error - }; - let scope_str: String = row.get("scope")?; let scope = match scope_str.as_str() { "session" => PermissionScope::Session, @@ -61,7 +48,7 @@ impl SqlitePermissionStore { tool_name: row.get("tool_name")?, command_pattern: row.get("command_pattern")?, resource_pattern: row.get("resource_pattern")?, - decision, + user_decision: UserPermissionDecision::AlwaysAllow, scope, project_id: row.get("project_id")?, created_at, @@ -71,12 +58,10 @@ impl SqlitePermissionStore { impl PermissionStore for SqlitePermissionStore { fn create_permission(&self, permission: Permission) -> Result { - let decision_str = match permission.decision { - PermissionDecision::AllowAllReadsInSession => "allow_all_reads_session", - PermissionDecision::AllowAllWritesInSession => "allow_all_writes_session", - PermissionDecision::AllowCommandForProject => "allow_command_for_project", - PermissionDecision::Ask => "ask", - PermissionDecision::AllowOnce => "once", + let decision_str = match permission.user_decision { + UserPermissionDecision::AlwaysAllow => "always_allow", + UserPermissionDecision::AllowOnce => "allow_once", + UserPermissionDecision::Deny => "deny", }; let scope_str = match permission.scope { @@ -159,10 +144,8 @@ impl PermissionStore for SqlitePermissionStore { })?; // Build params list dynamically to match the query - let mut param_values: Vec> = vec![ - Box::new(project_id), - Box::new(tool.to_string()), - ]; + let mut param_values: Vec> = + vec![Box::new(project_id), Box::new(tool.to_string())]; if !command_pattern.is_empty() { param_values.push(Box::new(command_pattern.to_string())); @@ -172,10 +155,8 @@ impl PermissionStore for SqlitePermissionStore { param_values.push(Box::new(resource_pattern.to_string())); } - let params_refs: Vec<&dyn rusqlite::ToSql> = param_values - .iter() - .map(|p| p.as_ref()) - .collect(); + let params_refs: Vec<&dyn rusqlite::ToSql> = + param_values.iter().map(|p| p.as_ref()).collect(); let result = conn.query_row(&query, params_refs.as_slice(), |row| { self.row_to_permission(row) @@ -187,32 +168,6 @@ impl PermissionStore for SqlitePermissionStore { Err(e) => Err(StoreError::Database(e)), } } - - fn find_session_permissions(&self, project_id: i32) -> Result, StoreError> { - let conn = self.conn.get().map_err(|e| { - StoreError::Database(rusqlite::Error::ToSqlConversionFailure(Box::new( - std::io::Error::new( - std::io::ErrorKind::Other, - format!("Failed to get connection: {}", e), - ), - ))) - })?; - - let mut stmt = conn.prepare( - "SELECT id, tool_name, command_pattern, resource_pattern, decision, scope, - project_id, created_at - FROM permissions - WHERE project_id = ?1 - AND scope = 'session' - AND decision IN ('allow_all_reads_session', 'allow_all_writes_session')", - )?; - - let permissions = stmt - .query_map(params![project_id], |row| self.row_to_permission(row))? - .collect::, _>>()?; - - Ok(permissions) - } } #[cfg(test)] @@ -238,7 +193,8 @@ mod tests { created_at TEXT NOT NULL )", [], - ).unwrap(); + ) + .unwrap(); pool } @@ -253,7 +209,7 @@ mod tests { "test_tool".to_string(), None, None, - PermissionDecision::AllowOnce, + UserPermissionDecision::AllowOnce, PermissionScope::Project, Some(1), ); @@ -262,7 +218,10 @@ mod tests { // Should find the permission when searching with empty strings let result = store.find_permission("test_tool", 1, "", "").unwrap(); assert!(result.is_some()); - assert_eq!(result.unwrap().decision, PermissionDecision::AllowOnce); + assert_eq!( + result.unwrap().user_decision, + UserPermissionDecision::AlwaysAllow + ); } #[test] @@ -275,16 +234,21 @@ mod tests { "test_tool".to_string(), Some("echo test".to_string()), None, - PermissionDecision::Ask, + UserPermissionDecision::AlwaysAllow, PermissionScope::Project, Some(1), ); store.create_permission(permission).unwrap(); // Should find the permission when searching with command and empty resource - let result = store.find_permission("test_tool", 1, "echo test", "").unwrap(); + let result = store + .find_permission("test_tool", 1, "echo test", "") + .unwrap(); assert!(result.is_some()); - assert_eq!(result.unwrap().decision, PermissionDecision::Ask); + assert_eq!( + result.unwrap().user_decision, + UserPermissionDecision::AlwaysAllow + ); } #[test] @@ -297,16 +261,21 @@ mod tests { "test_tool".to_string(), None, Some("/path/to/file".to_string()), - PermissionDecision::AllowOnce, + UserPermissionDecision::AllowOnce, PermissionScope::Project, Some(1), ); store.create_permission(permission).unwrap(); // Should find the permission when searching with empty command and resource - let result = store.find_permission("test_tool", 1, "", "/path/to/file").unwrap(); + let result = store + .find_permission("test_tool", 1, "", "/path/to/file") + .unwrap(); assert!(result.is_some()); - assert_eq!(result.unwrap().decision, PermissionDecision::AllowOnce); + assert_eq!( + result.unwrap().user_decision, + UserPermissionDecision::AlwaysAllow + ); } #[test] @@ -319,16 +288,21 @@ mod tests { "test_tool".to_string(), Some("rm -rf".to_string()), Some("/etc/passwd".to_string()), - PermissionDecision::Ask, + UserPermissionDecision::AlwaysAllow, PermissionScope::Project, Some(1), ); store.create_permission(permission).unwrap(); // Should find the permission when searching with both patterns - let result = store.find_permission("test_tool", 1, "rm -rf", "/etc/passwd").unwrap(); + let result = store + .find_permission("test_tool", 1, "rm -rf", "/etc/passwd") + .unwrap(); assert!(result.is_some()); - assert_eq!(result.unwrap().decision, PermissionDecision::Ask); + assert_eq!( + result.unwrap().user_decision, + UserPermissionDecision::AlwaysAllow + ); } #[test] @@ -341,7 +315,7 @@ mod tests { "test_tool".to_string(), Some("echo test".to_string()), None, - PermissionDecision::AllowOnce, + UserPermissionDecision::AllowOnce, PermissionScope::Project, Some(1), ); @@ -353,47 +327,23 @@ mod tests { } #[test] - fn test_find_session_permissions() { + fn test_find_permission_different_project_no_match() { let pool = setup_test_db(); let store = SqlitePermissionStore::new(pool.clone()); - // Create session-wide read permission - let read_perm = Permission::new( + // Create permission for project 1 + let permission = Permission::new( "read_tool".to_string(), None, - Some("/project/**".to_string()), - PermissionDecision::AllowAllReadsInSession, - PermissionScope::Session, - Some(1), - ); - store.create_permission(read_perm).unwrap(); - - // Create session-wide write permission - let write_perm = Permission::new( - "write_tool".to_string(), None, - Some("/project/**".to_string()), - PermissionDecision::AllowAllWritesInSession, - PermissionScope::Session, - Some(1), - ); - store.create_permission(write_perm).unwrap(); - - // Create non-session permission (should not be returned) - let other_perm = Permission::new( - "other_tool".to_string(), - None, - None, - PermissionDecision::AllowOnce, + UserPermissionDecision::AlwaysAllow, PermissionScope::Project, Some(1), ); - store.create_permission(other_perm).unwrap(); + store.create_permission(permission).unwrap(); - // Find session permissions - let result = store.find_session_permissions(1).unwrap(); - assert_eq!(result.len(), 2); - assert!(result.iter().any(|p| matches!(p.decision, PermissionDecision::AllowAllReadsInSession))); - assert!(result.iter().any(|p| matches!(p.decision, PermissionDecision::AllowAllWritesInSession))); + // Should NOT find when searching for project 2 + let result = store.find_permission("read_tool", 2, "", "").unwrap(); + assert!(result.is_none()); } -} \ No newline at end of file +} diff --git a/src/domain/permissions/types.rs b/src/domain/permissions/types.rs index 4014f8f..a798e34 100644 --- a/src/domain/permissions/types.rs +++ b/src/domain/permissions/types.rs @@ -3,12 +3,16 @@ use serde::{Deserialize, Serialize}; use std::path::PathBuf; #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] -pub enum PermissionDecision { - Ask, // Prompt user again - AllowOnce, // One-time approval - AllowAllReadsInSession, // Session-wide read access - AllowAllWritesInSession, // Session-wide write access - AllowCommandForProject, // Allow specific command for this project (persistent) +pub enum SystemPermissionDecision { + Ask, + Allow, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub enum UserPermissionDecision { + AllowOnce, + AlwaysAllow, + Deny, } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] @@ -24,9 +28,9 @@ pub struct Permission { pub tool_name: String, // "shell", "file_edit", etc. pub command_pattern: Option, // "rm -rf", "git push", etc. pub resource_pattern: Option, // "api.search.brave.com", "/src/**", etc. - pub decision: PermissionDecision, // Allow, Deny, Ask - pub scope: PermissionScope, // Session, Project, Global - pub project_id: Option, // If project-scoped + pub user_decision: UserPermissionDecision, + pub scope: PermissionScope, // Session, Project, Global + pub project_id: Option, // If project-scoped pub created_at: DateTime, } @@ -36,7 +40,7 @@ impl Permission { tool_name: String, command_pattern: Option, resource_pattern: Option, - decision: PermissionDecision, + decision: UserPermissionDecision, scope: PermissionScope, project_id: Option, ) -> Self { @@ -45,88 +49,18 @@ impl Permission { tool_name, command_pattern, resource_pattern, - decision, + user_decision: decision, scope, project_id, created_at: Utc::now(), } } - /// Check if this permission matches the given tool, command, and path - pub fn matches(&self, tool: &str, command: Option<&str>, path: Option<&PathBuf>) -> bool { - if self.tool_name != tool { - return false; + pub fn system_decision(self) -> SystemPermissionDecision { + if self.user_decision == UserPermissionDecision::AlwaysAllow { + return SystemPermissionDecision::Allow; } - - // Check command pattern if specified - if let (Some(pattern), Some(cmd)) = (&self.command_pattern, command) { - if !self.matches_command(pattern, cmd) { - return false; - } - } - - // Check resource pattern if specified - if let (Some(pattern), Some(p)) = (&self.resource_pattern, path) { - if !self.matches_path(pattern, p) { - return false; - } - } - - true - } - - fn matches_command(&self, pattern: &str, command: &str) -> bool { - // Simple pattern matching - can be enhanced with regex - if pattern.contains('*') { - // Basic wildcard matching - let pattern_parts: Vec<&str> = pattern.split_whitespace().collect(); - let command_parts: Vec<&str> = command.split_whitespace().collect(); - - if pattern_parts.len() != command_parts.len() { - return false; - } - - pattern_parts - .iter() - .zip(command_parts.iter()) - .all(|(p, c)| *p == "*" || *p == *c) - } else { - pattern == command - } - } - - fn matches_path(&self, pattern: &str, path: &PathBuf) -> bool { - let path_str = path.to_string_lossy(); - - if pattern.contains('*') { - // Basic glob matching - if pattern.ends_with("/**") { - let prefix = pattern.trim_end_matches("/**"); - path_str.starts_with(prefix) - } else if pattern.starts_with("**/") { - let suffix = pattern.trim_start_matches("**/"); - path_str.ends_with(suffix) - } else { - // Simple wildcard - can be enhanced with proper glob matching - path_str.contains(pattern.trim_matches('*')) - } - } else { - path_str == pattern - } - } - - /// Check if this is a session-wide "allow all reads" permission - pub fn is_session_wide_all_reads(&self) -> bool { - matches!(self.decision, PermissionDecision::AllowAllReadsInSession) - && self.scope == PermissionScope::Session - && self.resource_pattern.as_ref().map(|p| p.ends_with("/**")).unwrap_or(false) - } - - /// Check if this is a session-wide "allow all writes" permission - pub fn is_session_wide_all_writes(&self) -> bool { - matches!(self.decision, PermissionDecision::AllowAllWritesInSession) - && self.scope == PermissionScope::Session - && self.resource_pattern.as_ref().map(|p| p.ends_with("/**")).unwrap_or(false) + SystemPermissionDecision::Ask } } @@ -165,7 +99,7 @@ impl PermissionRequest { #[derive(Debug, Clone, Serialize, Deserialize)] pub struct PermissionConfig { - pub default_decision: PermissionDecision, + pub default_decision: SystemPermissionDecision, pub dangerous_commands: Vec, pub restricted_paths: Vec, pub require_confirmation: bool, @@ -174,7 +108,7 @@ pub struct PermissionConfig { impl Default for PermissionConfig { fn default() -> Self { Self { - default_decision: PermissionDecision::Ask, + default_decision: SystemPermissionDecision::Ask, dangerous_commands: vec![ "rm -rf".to_string(), "sudo".to_string(), @@ -188,6 +122,7 @@ impl Default for PermissionConfig { "/usr/bin".to_string(), "/bin".to_string(), "/sbin".to_string(), + ".env".to_string(), "~/.ssh".to_string(), "~/.aws".to_string(), "~/.gnupg".to_string(), @@ -195,4 +130,4 @@ impl Default for PermissionConfig { require_confirmation: true, } } -} \ No newline at end of file +} diff --git a/src/domain/prompting/general.rs b/src/domain/prompting/general.rs index 3e92dc2..a945f5f 100644 --- a/src/domain/prompting/general.rs +++ b/src/domain/prompting/general.rs @@ -1,8 +1,8 @@ -use crate::domain::AgentModeType; use crate::domain::bt::BTStepNodeInterface; -use crate::domain::session::Request; -use crate::domain::ModelType; use crate::domain::prompting::user::get_user_prompt; +use crate::domain::session::Request; +use crate::domain::AgentModeType; +use crate::domain::ModelType; /// LLM prompt templates for the coding assistant /// @@ -29,26 +29,36 @@ pub fn get_bt_tree_step_prompt( } } -pub fn get_system_prompt(model_type: ModelType, agent_mode: AgentModeType, remaining_calls: usize) -> String { +pub fn get_system_prompt( + model_type: ModelType, + agent_mode: AgentModeType, + remaining_calls: usize, +) -> String { let (os_name, shell_name) = get_runtime_environment(); if agent_mode == AgentModeType::Plan { let mut system_prompt = _system_prompt_for_plan(model_type); if remaining_calls < 3 { - system_prompt.push_str(&format!("\n\nYou have {} tool calls left to process this request.", remaining_calls)); + system_prompt.push_str(&format!( + "\n\nYou have {} tool calls left to process this request.", + remaining_calls + )); } return system_prompt; } if agent_mode == AgentModeType::BuildFromPlan { let mut system_prompt = _system_prompt_for_build_from_plan(model_type); if remaining_calls < 3 { - system_prompt.push_str(&format!("\n\nYou have {} tool calls left to process this request.", remaining_calls)); + system_prompt.push_str(&format!( + "\n\nYou have {} tool calls left to process this request.", + remaining_calls + )); } return system_prompt; } let mut system_prompt = if model_type == ModelType::OpenAI { format!( - "You are Drastis, a coding agent. \n\ + "You're Drastis, you're running as a coding agent in the CLI on a user's computer. \n\ Use the available tools to gather context and make changes. \ When using tools, pass JSON arguments that match their parameters. \n\ Runtime: os={}, shell={}.", @@ -56,12 +66,15 @@ pub fn get_system_prompt(model_type: ModelType, agent_mode: AgentModeType, remai ) } else { format!( - "You are Drastis, a coding agent. Use available tools to gather context and make changes. Be concise and accurate. Runtime: os={}, shell={}.", + "You're Drastis, you're running as a coding agent in the CLI on a user's computer. Use available tools to gather context and make changes. Be concise and accurate. Runtime: os={}, shell={}.", os_name, shell_name ) }; if remaining_calls < 3 { - system_prompt.push_str(&format!("\n\nYou have {} tool calls left to process this request.", remaining_calls)); + system_prompt.push_str(&format!( + "\n\nYou have {} tool calls left to process this request.", + remaining_calls + )); } system_prompt } @@ -70,7 +83,7 @@ fn _system_prompt_for_plan(model_type: ModelType) -> String { let (os_name, shell_name) = get_runtime_environment(); if model_type == ModelType::OpenAI { format!( - "You are Drastis, a coding agent. \n\ + "You're Drastis, you're running as a coding agent in the CLI on a user's computer. \n\ You're in the Plan Mode! Use tools to gather all required information to make a detailed plan of what user wants to achieve. \n\ Ask the user, if there're any ambiguities, and clarify what needs to be done if the instructions are not clear. \n\ When you gathered the information, use the `update_todo_list` tool to create the TODO list and ask the user if he would like to make any changes to the TODO list. \n\ diff --git a/src/domain/prompting/mod.rs b/src/domain/prompting/mod.rs index fbe4ab6..c2a3582 100644 --- a/src/domain/prompting/mod.rs +++ b/src/domain/prompting/mod.rs @@ -4,6 +4,6 @@ mod tools; mod user; pub use general::*; -pub use user::*; pub use session::*; pub use tools::*; +pub use user::*; diff --git a/src/domain/prompting/user.rs b/src/domain/prompting/user.rs index 8c0b80b..7cb3177 100644 --- a/src/domain/prompting/user.rs +++ b/src/domain/prompting/user.rs @@ -1,7 +1,7 @@ -use crate::domain::AgentModeType; -use crate::domain::ModelType; use crate::domain::session::Request; use crate::domain::todo::TodoListStatus; +use crate::domain::AgentModeType; +use crate::domain::ModelType; /// Format request history and current request as a complete prompt pub fn get_user_prompt(_model_type: ModelType, request: &dyn Request) -> String { @@ -23,4 +23,4 @@ fn _get_build_from_plan_prompt(request: &dyn Request) -> String { } } "Report very briefly what was done without using any tools, since all the TODO items are completed.".to_string() -} \ No newline at end of file +} diff --git a/src/domain/session/request.rs b/src/domain/session/request.rs index 0910558..b6aeb50 100644 --- a/src/domain/session/request.rs +++ b/src/domain/session/request.rs @@ -41,4 +41,4 @@ pub trait Request { /// Get the session's TODO list (plan) fn get_session_plan(&self) -> Option; -} \ No newline at end of file +} diff --git a/src/domain/session/service.rs b/src/domain/session/service.rs index 649fea9..fcadce1 100644 --- a/src/domain/session/service.rs +++ b/src/domain/session/service.rs @@ -7,8 +7,8 @@ use crate::infrastructure::db::DbPool; use crate::infrastructure::AgentToUiEvent; use crate::infrastructure::InferenceEngine; use crate::repository::{ - ModelsRepository, SessionRequestStepsRepository, SessionRequestsRepository, - TodoListRepository, UserSettingsRepository, + ModelsRepository, SessionRequestStepsRepository, SessionRequestsRepository, TodoListRepository, + UserSettingsRepository, }; use crossbeam_channel::Sender; use std::sync::Arc; @@ -97,7 +97,9 @@ impl SessionService { match repo.get_by_session(session.id()) { Ok(Some(row)) => { match serde_json::from_str::(&row.content) { - Ok(todo_list) => !todo_list.is_completed() && !todo_list.items.is_empty(), + Ok(todo_list) => { + !todo_list.is_completed() && !todo_list.items.is_empty() + } Err(_) => false, } } @@ -133,8 +135,7 @@ impl SessionService { // Run the workflow let result: Result = if self.use_behavior_trees { - self.workflow - .run_using_bt(&mut request, cancel) + self.workflow.run_using_bt(&mut request, cancel) } else { self.workflow .run(&mut request, cancel, effective_mode) diff --git a/src/domain/session/session.rs b/src/domain/session/session.rs index ec349e5..21361d1 100644 --- a/src/domain/session/session.rs +++ b/src/domain/session/session.rs @@ -128,6 +128,7 @@ impl Request for Session { } fn get_history_steps(&self) -> Vec { + use crate::domain::workflow::step::ChainStep; use crate::repository::SessionRequestStepsRepository; // BuildFromPlan keeps context minimal — TODO list (with statuses) is @@ -145,16 +146,43 @@ impl Request for Session { } }; - let steps_repo = SessionRequestStepsRepository::new(conn); + let steps_repo = SessionRequestStepsRepository::new(conn.clone()); // Load all steps from all requests in this session - match steps_repo.load_steps_for_session(self.id) { + let steps = match steps_repo.load_steps_for_session(self.id) { Ok(steps) => steps, Err(e) => { - log::warn!("Failed to load history steps for session {}: {}", self.id, e); + log::warn!( + "Failed to load history steps for session {}: {}", + self.id, + e + ); Vec::new() } + }; + + if !steps.is_empty() { + return steps; } + + // Fallback: rebuild assistant responses from stored steps_log when steps table is empty + let mut fallback_steps = Vec::new(); + for request in &self.requests { + if let Some(log) = request.steps_log() { + for line in log.lines() { + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + fallback_steps.push(ChainStep::assistant_response( + trimmed.to_string(), + trimmed.to_string(), + )); + } + } + } + + fallback_steps } fn get_session_plan(&self) -> Option { @@ -201,4 +229,4 @@ impl Request for Session { } // Note: From is no longer implemented because Session requires a Project entity. -// Use Session::from_row_with_project() instead. +// Use Session::from_row_with_project() instead. \ No newline at end of file diff --git a/src/domain/session/session_request.rs b/src/domain/session/session_request.rs index d0da4db..db4e50f 100644 --- a/src/domain/session/session_request.rs +++ b/src/domain/session/session_request.rs @@ -8,6 +8,7 @@ use crate::repository::SessionRequestRow; pub struct SessionRequest { prompt: String, result_summary: Option, + steps_log: Option, #[allow(dead_code)] mode: AgentModeType, } @@ -23,6 +24,10 @@ impl SessionRequest { self.result_summary.as_deref() } + pub fn steps_log(&self) -> Option<&str> { + self.steps_log.as_deref() + } + #[allow(dead_code)] pub fn mode(&self) -> AgentModeType { self.mode @@ -32,6 +37,7 @@ impl SessionRequest { Self { prompt: row.prompt, result_summary: row.result_summary, + steps_log: row._steps_log, mode: row.mode, } } diff --git a/src/domain/startup/service.rs b/src/domain/startup/service.rs index 53d5964..5b541e1 100644 --- a/src/domain/startup/service.rs +++ b/src/domain/startup/service.rs @@ -85,9 +85,10 @@ impl StartupService { let naming_prompt = session_naming_prompt(self.engine.get_type(), &prompt_preview); let mut chain = Chain::new(); - chain - .steps - .push(ChainStep::user_message(naming_prompt.parse().unwrap(), Vec::new())); + chain.steps.push(ChainStep::user_message( + naming_prompt.parse().unwrap(), + Vec::new(), + )); let session_name = match self.engine.generate(&[], &chain, &[], None) { Ok(raw) => { log::debug!("Raw session name response: {:?}", raw.summary); @@ -196,7 +197,8 @@ impl StartupService { .map(SessionRequest::from_row) .collect(); - let session = Session::load_with_requests(session_row, project, requests); + let mut session = Session::load_with_requests(session_row, project, requests); + session.set_conn(self.conn.clone()); Ok(session) } } diff --git a/src/domain/todo.rs b/src/domain/todo.rs index f01abcc..3418855 100644 --- a/src/domain/todo.rs +++ b/src/domain/todo.rs @@ -20,9 +20,15 @@ impl TodoListStatus { #[allow(dead_code)] pub fn from_str(value: &str) -> Self { match value { - "pending" | "PENDING" | "planned" | "PLANNED" | "todo" | "TODO" => TodoListStatus::Pending, - "in_progress" | "IN_PROGRESS" | "doing" | "DOING" | "started" | "STARTED" => TodoListStatus::InProgress, - "completed" | "COMPLETED" | "done" | "DONE" | "finished" | "FINISHED" => TodoListStatus::Completed, + "pending" | "PENDING" | "planned" | "PLANNED" | "todo" | "TODO" => { + TodoListStatus::Pending + } + "in_progress" | "IN_PROGRESS" | "doing" | "DOING" | "started" | "STARTED" => { + TodoListStatus::InProgress + } + "completed" | "COMPLETED" | "done" | "DONE" | "finished" | "FINISHED" => { + TodoListStatus::Completed + } &_ => TodoListStatus::Pending, } } @@ -42,6 +48,8 @@ pub struct TodoList { impl TodoList { pub fn is_completed(&self) -> bool { - self.items.iter().all(|item| item.status == TodoListStatus::Completed) + self.items + .iter() + .all(|item| item.status == TodoListStatus::Completed) } -} \ No newline at end of file +} diff --git a/src/domain/tools/discover_objects.rs b/src/domain/tools/discover_objects.rs index f06664c..be9beef 100644 --- a/src/domain/tools/discover_objects.rs +++ b/src/domain/tools/discover_objects.rs @@ -106,7 +106,12 @@ impl Tool for DiscoverObjects { let input = match self.load_input() { Ok(input) => input, Err(e) => { - return ToolResult::error(self.name().to_string(), String::new(), e.to_string(), String::new()) + return ToolResult::error( + self.name().to_string(), + String::new(), + e.to_string(), + String::new(), + ) } }; @@ -117,7 +122,12 @@ impl Tool for DiscoverObjects { Self::format_output(lang, &objects), input.call_id, ), - Err(e) => ToolResult::error(self.name().to_string(), input.raw.clone(), e.to_string(), input.call_id), + Err(e) => ToolResult::error( + self.name().to_string(), + input.raw.clone(), + e.to_string(), + input.call_id, + ), } } diff --git a/src/domain/tools/find_files.rs b/src/domain/tools/find_files.rs index b033da3..9c87e14 100644 --- a/src/domain/tools/find_files.rs +++ b/src/domain/tools/find_files.rs @@ -169,7 +169,12 @@ impl Tool for FindFiles { let input = match self.load_input() { Ok(input) => input, Err(e) => { - return ToolResult::error(self.name().to_string(), String::new(), e.to_string(), String::new()) + return ToolResult::error( + self.name().to_string(), + String::new(), + e.to_string(), + String::new(), + ) } }; @@ -177,7 +182,14 @@ impl Tool for FindFiles { let search_root = match Self::resolve_search_root(input.root.as_deref(), request.project_root()) { Ok(root) => root, - Err(e) => return ToolResult::error(self.name().to_string(), input.raw.clone(), e, input.call_id.clone()), + Err(e) => { + return ToolResult::error( + self.name().to_string(), + input.raw.clone(), + e, + input.call_id.clone(), + ) + } }; // Check if search root exists diff --git a/src/domain/tools/mod.rs b/src/domain/tools/mod.rs index cc75ea9..272f09b 100644 --- a/src/domain/tools/mod.rs +++ b/src/domain/tools/mod.rs @@ -47,7 +47,7 @@ pub struct ToolResult { tool_name: String, call_id: String, pub(crate) input: String, -is_successful: bool, + is_successful: bool, output: String, error_message: String, file_changes: Option>, @@ -62,7 +62,7 @@ impl ToolResult { output, error_message: String::new(), file_changes: None, - call_id + call_id, } } @@ -74,7 +74,7 @@ impl ToolResult { output: String::new(), error_message: message.into(), file_changes: None, - call_id + call_id, } } @@ -161,6 +161,7 @@ pub trait Tool: Send + Sync { fn is_read_only(&self) -> bool { false } + #[allow(dead_code)] fn skip_permission_check(&self) -> bool { false } diff --git a/src/domain/tools/patch_files.rs b/src/domain/tools/patch_files.rs index ff138df..6d85664 100644 --- a/src/domain/tools/patch_files.rs +++ b/src/domain/tools/patch_files.rs @@ -78,7 +78,12 @@ impl Tool for PatchFiles { let input = match self.load_input() { Ok(input) => input, Err(e) => { - return ToolResult::error(self.name().to_string(), String::new(), e.to_string(), String::new()) + return ToolResult::error( + self.name().to_string(), + String::new(), + e.to_string(), + String::new(), + ) } }; @@ -149,7 +154,7 @@ impl Tool for PatchFiles { return ToolResult::error( self.name().to_string(), input.raw.clone(), - format!("Failed to apply patch: {}", e), + format!("Patch Failed: {}. Next step: read the file region you’re editing (function/block + ~10 lines around it) and generate a new patch using those exact lines as context. Then retry patch_files.", e), input.call_id.clone(), ) } @@ -612,4 +617,4 @@ mod tests { result.output_string() ); } -} \ No newline at end of file +} diff --git a/src/domain/tools/read_objects.rs b/src/domain/tools/read_objects.rs index ef9e66f..6918b7f 100644 --- a/src/domain/tools/read_objects.rs +++ b/src/domain/tools/read_objects.rs @@ -202,7 +202,12 @@ impl Tool for ReadObjects { let input = match self.load_input() { Ok(input) => input, Err(e) => { - return ToolResult::error(self.name().to_string(), String::new(), e.to_string(), String::new()) + return ToolResult::error( + self.name().to_string(), + String::new(), + e.to_string(), + String::new(), + ) } }; @@ -213,7 +218,12 @@ impl Tool for ReadObjects { Self::format_output(lang, results), input.call_id, ), - Err(e) => ToolResult::error(self.name().to_string(), input.raw.clone(), e.to_string(), input.call_id), + Err(e) => ToolResult::error( + self.name().to_string(), + input.raw.clone(), + e.to_string(), + input.call_id, + ), } } @@ -485,4 +495,4 @@ mod tests { }; assert!(query.matches(&obj_partial), "Should match partial name"); } -} \ No newline at end of file +} diff --git a/src/domain/tools/shell_exec.rs b/src/domain/tools/shell_exec.rs index 387e7ff..a1ec056 100644 --- a/src/domain/tools/shell_exec.rs +++ b/src/domain/tools/shell_exec.rs @@ -28,6 +28,35 @@ struct ShellExecInputParsed { call_id: String, } +fn is_allowed_read_only_command(command: &str) -> bool { + let trimmed = command.trim(); + if trimmed.is_empty() { + return false; + } + + if trimmed.contains(';') + || trimmed.contains("&&") + || trimmed.contains("||") + || trimmed.contains('|') + || trimmed.contains('>') + || trimmed.contains('<') + { + return false; + } + + let mut parts = trimmed.split_whitespace(); + let first = parts.next().unwrap_or_default(); + let args: Vec<&str> = parts.collect(); + + match first { + "rg" | "grep" | "glob" | "cat" | "head" | "tail" | "less" | "more" | "wc" | "cut" + | "sort" | "uniq" | "find" | "ls" | "tree" | "stat" | "file" | "awk" | "pwd" + | "which" | "type" => true, + "sed" => args.iter().any(|arg| *arg == "-n" || *arg == "--quiet" || *arg == "--silent"), + _ => false, + } +} + impl Tool for ShellExec { fn name(&self) -> &'static str { "shell_exec" @@ -135,7 +164,12 @@ impl Tool for ShellExec { if !stderr.is_empty() { result.push_str(&format!("[stderr]: {}", stderr)); } - return ToolResult::error(self.name().to_string(), input.raw.clone(), result, input.call_id.clone()); + return ToolResult::error( + self.name().to_string(), + input.raw.clone(), + result, + input.call_id.clone(), + ); }; ToolResult::ok(self.name().to_string(), input.raw, result, input.call_id) @@ -251,6 +285,15 @@ DO NOT use it to change files, use `patch_files` tool for this.", paths } + + fn is_read_only(&self) -> bool { + self.input + .lock() + .unwrap() + .as_ref() + .map(|input| is_allowed_read_only_command(&input.command)) + .unwrap_or(false) + } } impl ShellExec { @@ -346,7 +389,10 @@ mod tests { let tool = ShellExec::new(); assert!(tool - .parse_input(r#"{"command":"echo hello"}"#.to_string(), "call-id".to_string()) + .parse_input( + r#"{"command":"echo hello"}"#.to_string(), + "call-id".to_string() + ) .is_none()); let result = tool.work(&request); @@ -388,10 +434,13 @@ mod tests { let tool = ShellExec::new(); assert!(tool - .parse_input(format!( - r#"{{"command":"pwd","working_dir":"{}"}}"#, - subdir.display() - ), "call-id".to_string()) + .parse_input( + format!( + r#"{{"command":"pwd","working_dir":"{}"}}"#, + subdir.display() + ), + "call-id".to_string() + ) .is_none()); let result = tool.work(&request); @@ -412,7 +461,10 @@ mod tests { let tool = ShellExec::new(); assert!(tool - .parse_input(r#"{"command":"pwd","working_dir":"/tmp"}"#.to_string(), "call-id".to_string()) + .parse_input( + r#"{"command":"pwd","working_dir":"/tmp"}"#.to_string(), + "call-id".to_string() + ) .is_none()); let result = tool.work(&request); @@ -450,7 +502,10 @@ mod tests { let tool = ShellExec::new(); assert!(tool - .parse_input(r#"{"command":"echo error >&2 && exit 1"}"#.to_string(), "call-id".to_string()) + .parse_input( + r#"{"command":"echo error >&2 && exit 1"}"#.to_string(), + "call-id".to_string() + ) .is_none()); let result = tool.work(&request); @@ -478,10 +533,13 @@ mod tests { let tool = ShellExec::new(); assert!(tool - .parse_input(format!( - r#"{{"command":"echo 'test content' > {}"}}"#, - file_path.display() - ), "call-id".to_string()) + .parse_input( + format!( + r#"{{"command":"echo 'test content' > {}"}}"#, + file_path.display() + ), + "call-id".to_string() + ) .is_none()); let result = tool.work(&request); @@ -503,7 +561,10 @@ mod tests { let tool = ShellExec::new(); assert!(tool - .parse_input(r#"{"command":"echo 'hello world' | tr 'a-z' 'A-Z'"}"#.to_string(), "call-id".to_string()) + .parse_input( + r#"{"command":"echo 'hello world' | tr 'a-z' 'A-Z'"}"#.to_string(), + "call-id".to_string() + ) .is_none()); let result = tool.work(&request); @@ -513,4 +574,64 @@ mod tests { result.output_string() ); } + + #[test] + fn test_shell_exec_is_read_only_allows_read_commands() { + let tool = ShellExec::new(); + assert!(tool + .parse_input( + r#"{"command":"rg -n foo src"}"#.to_string(), + "call-id".to_string() + ) + .is_none()); + assert!(tool.is_read_only()); + + let tool = ShellExec::new(); + assert!(tool + .parse_input( + r#"{"command":"grep -n foo file.txt"}"#.to_string(), + "call-id".to_string() + ) + .is_none()); + assert!(tool.is_read_only()); + + let tool = ShellExec::new(); + assert!(tool + .parse_input( + r#"{"command":"sed -n '1,10p' file.txt"}"#.to_string(), + "call-id".to_string() + ) + .is_none()); + assert!(tool.is_read_only()); + } + + #[test] + fn test_shell_exec_is_read_only_rejects_write_or_compound() { + let tool = ShellExec::new(); + assert!(tool + .parse_input( + r#"{"command":"sed '1,10p' file.txt"}"#.to_string(), + "call-id".to_string() + ) + .is_none()); + assert!(!tool.is_read_only()); + + let tool = ShellExec::new(); + assert!(tool + .parse_input( + r#"{"command":"echo hi > out.txt"}"#.to_string(), + "call-id".to_string() + ) + .is_none()); + assert!(!tool.is_read_only()); + + let tool = ShellExec::new(); + assert!(tool + .parse_input( + r#"{"command":"echo hi && rg -n foo src"}"#.to_string(), + "call-id".to_string() + ) + .is_none()); + assert!(!tool.is_read_only()); + } } \ No newline at end of file diff --git a/src/domain/tools/structure.rs b/src/domain/tools/structure.rs index 82f6c16..b8cbb0c 100644 --- a/src/domain/tools/structure.rs +++ b/src/domain/tools/structure.rs @@ -164,13 +164,25 @@ impl Tool for Structure { let input = match self.load_input() { Ok(input) => input, Err(e) => { - return ToolResult::error(self.name().to_string(), String::new(), e.to_string(), String::new()) + return ToolResult::error( + self.name().to_string(), + String::new(), + e.to_string(), + String::new(), + ) } }; let target_path = match Self::resolve_path(&input.path, request.project_root()) { Ok(p) => p, - Err(e) => return ToolResult::error(self.name().to_string(), input.raw.clone(), e, input.call_id.clone()), + Err(e) => { + return ToolResult::error( + self.name().to_string(), + input.raw.clone(), + e, + input.call_id.clone(), + ) + } }; if !target_path.exists() { diff --git a/src/domain/tools/update_todo_list.rs b/src/domain/tools/update_todo_list.rs index 6c2e6d1..f090baa 100644 --- a/src/domain/tools/update_todo_list.rs +++ b/src/domain/tools/update_todo_list.rs @@ -1,4 +1,5 @@ use crate::domain::session::Request; +use crate::domain::todo::TodoList; use crate::domain::tools::{Error, Tool, ToolResult}; use crate::infrastructure::db::DbPool; use crate::infrastructure::event_bus::AgentToUiEvent; @@ -6,7 +7,6 @@ use crate::repository::TodoListRepository; use crossbeam_channel::Sender; use serde_json::{json, Value}; use std::sync::Mutex; -use crate::domain::todo::TodoList; pub struct UpdateTodoList { input: Mutex>, @@ -42,7 +42,10 @@ impl Tool for UpdateTodoList { match parsed { Ok(_) => { let mut lock = self.input.lock().unwrap(); - *lock = Some(UpdateTodoListInput { raw: input, call_id }); + *lock = Some(UpdateTodoListInput { + raw: input, + call_id, + }); None } Err(e) => Some(Error::Parse(format!("Failed to parse input: {}", e))), @@ -148,7 +151,7 @@ impl Tool for UpdateTodoList { } fn desc(&self) -> String { -"Use this in Plan mode to create or update TODO items after exploring the codebase. For every use, provide full TODO list, as it will replace the existing list.\n\ + "Use this in Plan mode to create or update TODO items after exploring the codebase. For every use, provide full TODO list, as it will replace the existing list.\n\ If you're not in Plan mode, only update statuses of TODO items as you finish them. " .to_string() diff --git a/src/domain/tools/web_search.rs b/src/domain/tools/web_search.rs index 8444d76..3cfe84e 100644 --- a/src/domain/tools/web_search.rs +++ b/src/domain/tools/web_search.rs @@ -163,9 +163,12 @@ impl Tool for WebSearch { Ok(results) => { let output = WebSearchOutput { results }; match serde_json::to_string(&output) { - Ok(json_output) => { - ToolResult::ok(self.name().to_string(), input.raw.clone(), json_output, input.call_id) - } + Ok(json_output) => ToolResult::ok( + self.name().to_string(), + input.raw.clone(), + json_output, + input.call_id, + ), Err(e) => ToolResult::error( self.name().to_string(), input.raw.clone(), @@ -320,7 +323,10 @@ mod tests { use crate::domain::permissions::{ checker::{PermissionChecker, PermissionPrompter}, store::{PermissionStore, StoreError}, - types::{PermissionConfig, PermissionDecision, PermissionRequest, PermissionScope}, + types::{ + PermissionConfig, PermissionRequest, PermissionScope, SystemPermissionDecision, + UserPermissionDecision, + }, }; use crate::domain::session::{Request, SessionRequest}; use crate::domain::UserSettings; @@ -359,11 +365,6 @@ mod tests { Ok(permission) } - fn find_session_permissions(&self, _project_id: i32) -> Result, StoreError> { - // Return empty for tests - this test doesn't use session permissions - Ok(Vec::new()) - } - fn find_permission( &self, tool: &str, @@ -407,14 +408,14 @@ mod tests { struct TestPrompter { calls: Arc, - decision: PermissionDecision, + decision: UserPermissionDecision, } impl PermissionPrompter for TestPrompter { fn ask_permission( &self, _request: &PermissionRequest, - ) -> Result { + ) -> Result { self.calls.fetch_add(1, Ordering::SeqCst); Ok(self.decision.clone()) } @@ -482,6 +483,11 @@ mod tests { web_search_enabled, brave_api_key: brave_api_key.map(String::from), max_tool_calls_per_request: 10, + auth_method: "api_key".to_string(), + oauth_access_token: None, + oauth_refresh_token: None, + oauth_token_expiry: None, + oauth_account_id: None, }; UserSettings::from(row) } @@ -495,12 +501,12 @@ mod tests { let calls = Arc::new(AtomicUsize::new(0)); let prompter = Arc::new(TestPrompter { calls: Arc::clone(&calls), - decision: PermissionDecision::AllowOnce, + decision: UserPermissionDecision::AllowOnce, }); let checker = PermissionChecker::new_with_prompter( store, PermissionConfig { - default_decision: PermissionDecision::Ask, + default_decision: SystemPermissionDecision::Ask, ..PermissionConfig::default() }, prompter, @@ -525,12 +531,12 @@ mod tests { let store = Arc::new(TestStore::new()); let calls = Arc::new(AtomicUsize::new(0)); - let decision_to_return = Arc::new(Mutex::new(PermissionDecision::AllowOnce)); + let decision_to_return = Arc::new(Mutex::new(UserPermissionDecision::AllowOnce)); let captured_request = Arc::new(Mutex::new(None::)); struct CapturingPrompter { calls: Arc, - decision: Arc>, + decision: Arc>, captured: Arc>>, } @@ -538,7 +544,7 @@ mod tests { fn ask_permission( &self, request: &PermissionRequest, - ) -> Result { + ) -> Result { self.calls.fetch_add(1, Ordering::SeqCst); *self.captured.lock().unwrap() = Some(request.clone()); Ok(self.decision.lock().unwrap().clone()) @@ -554,7 +560,7 @@ mod tests { let checker = PermissionChecker::new_with_prompter( store, PermissionConfig { - default_decision: PermissionDecision::Ask, + default_decision: SystemPermissionDecision::Ask, ..PermissionConfig::default() }, prompter, @@ -578,5 +584,4 @@ mod tests { assert_eq!(captured_req.tool_name, "web_search"); assert_eq!(captured_req.paths, vec![PathBuf::from("docs.rs")]); } - -} \ No newline at end of file +} diff --git a/src/domain/user_settings.rs b/src/domain/user_settings.rs index 8d2c4b1..c9ac97d 100644 --- a/src/domain/user_settings.rs +++ b/src/domain/user_settings.rs @@ -1,6 +1,30 @@ use crate::repository::UserSettingsRow; +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum AuthMethod { + ApiKey, + OAuth, +} + +#[allow(dead_code)] +impl AuthMethod { + pub fn as_str(&self) -> &str { + match self { + AuthMethod::ApiKey => "api_key", + AuthMethod::OAuth => "oauth", + } + } + + pub fn from_str(s: &str) -> Self { + match s { + "oauth" => AuthMethod::OAuth, + _ => AuthMethod::ApiKey, + } + } +} + #[derive(Debug, Clone)] +#[allow(dead_code)] pub struct UserSettings { _id: i64, openai_api_key: Option, @@ -11,8 +35,14 @@ pub struct UserSettings { web_search_enabled: bool, brave_api_key: Option, max_tool_calls_per_request: i32, + auth_method: AuthMethod, + oauth_access_token: Option, + oauth_refresh_token: Option, + oauth_token_expiry: Option, + oauth_account_id: Option, } +#[allow(dead_code)] impl UserSettings { pub fn openai_api_key(&self) -> Option<&str> { self.openai_api_key.as_deref() @@ -22,7 +52,6 @@ impl UserSettings { self.openai_tracing_enabled } - #[allow(dead_code)] pub fn current_model_name(&self) -> Option<&str> { self.current_model_name.as_deref() } @@ -43,6 +72,29 @@ impl UserSettings { pub fn max_tool_calls_per_request(&self) -> i32 { self.max_tool_calls_per_request } + + pub fn auth_method(&self) -> &AuthMethod { + &self.auth_method + } + + pub fn oauth_access_token(&self) -> Option<&str> { + self.oauth_access_token.as_deref() + } + + pub fn oauth_refresh_token(&self) -> Option<&str> { + self.oauth_refresh_token.as_deref() + } + + pub fn oauth_account_id(&self) -> Option<&str> { + self.oauth_account_id.as_deref() + } + + pub fn is_oauth_token_expired(&self) -> bool { + if let Some(expiry) = self.oauth_token_expiry { + return chrono::Utc::now().timestamp() >= expiry; + } + false + } } impl From for UserSettings { @@ -57,6 +109,11 @@ impl From for UserSettings { web_search_enabled: row.web_search_enabled, brave_api_key: row.brave_api_key, max_tool_calls_per_request: row.max_tool_calls_per_request, + auth_method: AuthMethod::from_str(&row.auth_method), + oauth_access_token: row.oauth_access_token, + oauth_refresh_token: row.oauth_refresh_token, + oauth_token_expiry: row.oauth_token_expiry, + oauth_account_id: row.oauth_account_id, } } } diff --git a/src/domain/workflow/chain.rs b/src/domain/workflow/chain.rs index be07c0a..659815d 100644 --- a/src/domain/workflow/chain.rs +++ b/src/domain/workflow/chain.rs @@ -1,8 +1,8 @@ use super::step::ChainStep; +use super::step::StepType; +use crate::domain::todo::TodoList; use crate::domain::tools::ToolResult; use serde::{Deserialize, Serialize}; -use crate::domain::todo::TodoList; -use super::step::StepType; /// Represents an execution chain containing multiple steps /// The chain is built incrementally as tools are executed @@ -36,7 +36,8 @@ impl Chain { /// Add a step to the chain after executing a tool pub fn add_step(&mut self, result: ToolResult) { if result.tool_name() == "update_todo_list" && result.is_successful() { - if let Ok(updated_todo_list) = serde_json::from_str::(&result.input_string()) { + if let Ok(updated_todo_list) = serde_json::from_str::(&result.input_string()) + { self.set_todo_list(Some(updated_todo_list)); return; } @@ -122,7 +123,7 @@ impl Chain { .collect::>() .join("\n") } - + pub fn set_system_prompt(&mut self, system_prompt: String) { self.system_prompt = system_prompt; } diff --git a/src/domain/workflow/toolset/all.rs b/src/domain/workflow/toolset/all.rs index 4d9a124..e7810db 100644 --- a/src/domain/workflow/toolset/all.rs +++ b/src/domain/workflow/toolset/all.rs @@ -16,7 +16,7 @@ impl AllToolset { session_id: i64, settings: &UserSettings, conn: DbPool, - event_sender: Sender + event_sender: Sender, ) -> Self { let mut tools: HashMap> = HashMap::new(); diff --git a/src/domain/workflow/toolset/all_no_todo.rs b/src/domain/workflow/toolset/all_no_todo.rs index cb53bf3..0e11ec6 100644 --- a/src/domain/workflow/toolset/all_no_todo.rs +++ b/src/domain/workflow/toolset/all_no_todo.rs @@ -15,7 +15,7 @@ impl AllNoTodoToolset { _session_id: i64, settings: &UserSettings, _conn: DbPool, - _event_sender: Sender + _event_sender: Sender, ) -> Self { let mut tools: HashMap> = HashMap::new(); @@ -50,4 +50,4 @@ impl Toolset for AllNoTodoToolset { fn tools(&self) -> &HashMap> { &self.tools } -} \ No newline at end of file +} diff --git a/src/domain/workflow/toolset/discover.rs b/src/domain/workflow/toolset/discover.rs index e3fea00..367448f 100644 --- a/src/domain/workflow/toolset/discover.rs +++ b/src/domain/workflow/toolset/discover.rs @@ -1,9 +1,9 @@ use crate::domain::tools::*; use crate::domain::workflow::toolset::Toolset; use crate::infrastructure::db::DbPool; -use std::collections::HashMap; -use crossbeam_channel::Sender; use crate::infrastructure::AgentToUiEvent; +use crossbeam_channel::Sender; +use std::collections::HashMap; /// General toolset containing read-only and utility tools pub struct DiscoverToolset { diff --git a/src/domain/workflow/toolset/finishing_all.rs b/src/domain/workflow/toolset/finishing_all.rs index a7765e6..5cc3c85 100644 --- a/src/domain/workflow/toolset/finishing_all.rs +++ b/src/domain/workflow/toolset/finishing_all.rs @@ -11,11 +11,7 @@ pub struct FinishingAllToolset { } impl FinishingAllToolset { - pub fn new( - session_id: i64, - conn: DbPool, - event_sender: Sender - ) -> Self { + pub fn new(session_id: i64, conn: DbPool, event_sender: Sender) -> Self { let mut tools: HashMap> = HashMap::new(); let patch_files = Box::new(PatchFiles::new()); diff --git a/src/domain/workflow/toolset/finishing_all_no_todo.rs b/src/domain/workflow/toolset/finishing_all_no_todo.rs index d06bdfc..59fce00 100644 --- a/src/domain/workflow/toolset/finishing_all_no_todo.rs +++ b/src/domain/workflow/toolset/finishing_all_no_todo.rs @@ -10,11 +10,7 @@ pub struct FinishingAllNoTodoToolset { } impl FinishingAllNoTodoToolset { - pub fn new( - _session_id: i64, - _conn: DbPool, - _event_sender: Sender - ) -> Self { + pub fn new(_session_id: i64, _conn: DbPool, _event_sender: Sender) -> Self { let mut tools: HashMap> = HashMap::new(); let patch_files = Box::new(PatchFiles::new()); @@ -28,4 +24,4 @@ impl Toolset for FinishingAllNoTodoToolset { fn tools(&self) -> &HashMap> { &self.tools } -} \ No newline at end of file +} diff --git a/src/domain/workflow/toolset/mod.rs b/src/domain/workflow/toolset/mod.rs index 0131dee..a84958e 100644 --- a/src/domain/workflow/toolset/mod.rs +++ b/src/domain/workflow/toolset/mod.rs @@ -1,29 +1,29 @@ mod all; mod all_no_todo; -mod edit; -mod none; mod discover; +mod edit; mod finishing_all; mod finishing_all_no_todo; mod finishing_discover; +mod none; -use std::collections::HashMap; -use std::sync::Arc; -use crossbeam_channel::Sender; +use crate::domain::tools::Error; +use crate::domain::tools::Tool; +use crate::domain::UserSettings; +use crate::infrastructure::db::DbPool; +use crate::infrastructure::event_bus::AgentToUiEvent; +use crate::infrastructure::inference::ToolCall; pub use all::AllToolset; pub use all_no_todo::AllNoTodoToolset; -pub use edit::EditToolset; -pub use none::NoneToolset; +use crossbeam_channel::Sender; pub use discover::DiscoverToolset; +pub use edit::EditToolset; pub use finishing_all::FinishingAllToolset; pub use finishing_all_no_todo::FinishingAllNoTodoToolset; pub use finishing_discover::FinishingDiscoverToolset; -use crate::infrastructure::db::DbPool; -use crate::domain::tools::Tool; -use crate::domain::UserSettings; -use crate::infrastructure::event_bus::AgentToUiEvent; -use crate::infrastructure::inference::ToolCall; -use crate::domain::tools::Error; +pub use none::NoneToolset; +use std::collections::HashMap; +use std::sync::Arc; #[derive(Clone, Copy)] pub enum ToolsetType { @@ -51,12 +51,29 @@ impl ToolsetType { ToolsetType::Discover => Arc::new(DiscoverToolset::new(session_id, conn, event_sender)), ToolsetType::Edit => Arc::new(EditToolset::new(session_id, conn, event_sender)), ToolsetType::All => Arc::new(AllToolset::new(session_id, settings, conn, event_sender)), - ToolsetType::AllNoTodo => Arc::new(AllNoTodoToolset::new(session_id, settings, conn, event_sender)), - ToolsetType::FinishingAllNoTodo => Arc::new(FinishingAllNoTodoToolset::new(session_id, conn, event_sender)), + ToolsetType::AllNoTodo => Arc::new(AllNoTodoToolset::new( + session_id, + settings, + conn, + event_sender, + )), + ToolsetType::FinishingAllNoTodo => Arc::new(FinishingAllNoTodoToolset::new( + session_id, + conn, + event_sender, + )), ToolsetType::None => Arc::new(NoneToolset::new()), - ToolsetType::FinishingAll => Arc::new(FinishingAllToolset::new(session_id, conn, event_sender)), - ToolsetType::FinishingDiscover => Arc::new(FinishingDiscoverToolset::new(session_id, conn, event_sender)), - ToolsetType::FinishingEdit => Arc::new(EditToolset::new(session_id, conn, event_sender)), + ToolsetType::FinishingAll => { + Arc::new(FinishingAllToolset::new(session_id, conn, event_sender)) + } + ToolsetType::FinishingDiscover => Arc::new(FinishingDiscoverToolset::new( + session_id, + conn, + event_sender, + )), + ToolsetType::FinishingEdit => { + Arc::new(EditToolset::new(session_id, conn, event_sender)) + } } } @@ -102,10 +119,11 @@ pub trait Toolset { .map(|t| t.as_ref()) .ok_or_else(|| Error::Parse(format!("Tool not found: {}", tool_call.name)))?; - if let Some(err) = tool.parse_input(tool_call.arguments.clone(), tool_call.call_id.clone()) { + if let Some(err) = tool.parse_input(tool_call.arguments.clone(), tool_call.call_id.clone()) + { return Err(err); } Ok(tool) } -} \ No newline at end of file +} diff --git a/src/domain/workflow/workflow.rs b/src/domain/workflow/workflow.rs index c9c188b..7da4aec 100644 --- a/src/domain/workflow/workflow.rs +++ b/src/domain/workflow/workflow.rs @@ -6,17 +6,17 @@ use crate::domain::bt::GeneralTree; use crate::domain::prompting; use crate::domain::session::Request; use crate::domain::todo::TodoListStatus; +use crate::domain::workflow::toolset::NoneToolset; +use crate::domain::workflow::toolset::ToolsetType; +use crate::domain::AgentModeType; use crate::infrastructure::db::DbPool; use crate::infrastructure::event_bus::{AgentToUiEvent, StepPhase}; use crate::infrastructure::inference::InferenceEngine; -use tokio::runtime::{Builder, Handle}; -use crate::domain::workflow::toolset::NoneToolset; -use crate::domain::workflow::toolset::ToolsetType; +use crossbeam_channel::Sender; use openai_agents_tracing::{SpanKind, TracingFacade}; use std::sync::Arc; use std::time::{SystemTime, UNIX_EPOCH}; -use crossbeam_channel::Sender; -use crate::domain::AgentModeType; +use tokio::runtime::{Builder, Handle}; /// Main workflow orchestrator that runs LLM-driven coding tasks /// Implements an eternal agent loop that: @@ -48,10 +48,7 @@ impl Workflow { toolset: Arc::new(NoneToolset::new()), engine, chain: Chain::new(), - tool_runner: ToolRunner::new( - permission_checker, - event_sender.clone(), - ), + tool_runner: ToolRunner::new(permission_checker, event_sender.clone()), tracer: None, event_sender, conn, @@ -170,17 +167,20 @@ impl Workflow { user_prompt_override: Option, mode: AgentModeType, ) -> Result<(), Error> { - let base_user_prompt = prompting::get_user_prompt(self.engine.get_type(), request); let user_prompt = user_prompt_override .as_deref() .unwrap_or(&base_user_prompt) .to_string(); - self.chain.steps.push(ChainStep::user_message(user_prompt.clone(), request.images().to_vec())); + self.chain.steps.push(ChainStep::user_message( + user_prompt.clone(), + request.images().to_vec(), + )); // Get max_tool_calls from override (for BT mode) or user settings let max_tool_calls = max_tool_calls_override.unwrap_or_else(|| { - request.user_settings() + request + .user_settings() .map(|s| s.max_tool_calls_per_request() as usize) .unwrap_or(50) }); @@ -189,7 +189,11 @@ impl Workflow { let mut tool_call_count = 0; let mut current_active_todo = self.chain.todo_list.clone(); let mut in_finishing_mode = false; - let finishing_threshold = if max_tool_calls > 5 { max_tool_calls - 5 } else { max_tool_calls }; + let finishing_threshold = if max_tool_calls > 5 { + max_tool_calls - 5 + } else { + max_tool_calls + }; // Eternal agent loop for _iteration in 1..=max_tool_calls { @@ -204,7 +208,11 @@ impl Workflow { let remaining_calls = max_tool_calls.saturating_sub(tool_call_count); // Get base system prompt and inject remaining count - let system_prompt = prompting::get_system_prompt(self.engine.get_type(), request.mode(), remaining_calls); + let system_prompt = prompting::get_system_prompt( + self.engine.get_type(), + request.mode(), + remaining_calls, + ); self.chain.set_system_prompt(system_prompt.clone()); // Switch to finishing toolset if approaching limit @@ -280,7 +288,10 @@ impl Workflow { // Check if TODO item changed and reset counter if so (BuildFromPlan mode only) if mode == AgentModeType::BuildFromPlan && is_update_todo { if self._did_todo_item_change(¤t_active_todo) { - log::info!("Active TODO item changed, resetting counter from {} to 0", tool_call_count); + log::info!( + "Active TODO item changed, resetting counter from {} to 0", + tool_call_count + ); tool_call_count = 0; current_active_todo = self.chain.todo_list.clone(); @@ -364,29 +375,28 @@ impl Workflow { .map(|duration| duration.as_nanos()) .unwrap_or(0); let index = (nanos % options.len() as u128) as usize; - let _ = self - .event_sender - .send(AgentToUiEvent::ProgressEvent { - step_name: "inference".to_string(), - phase: StepPhase::Before, - summary: options[index].to_string(), - }); + let _ = self.event_sender.send(AgentToUiEvent::ProgressEvent { + step_name: "inference".to_string(), + phase: StepPhase::Before, + summary: options[index].to_string(), + }); } /// Detects if the active TODO item has changed /// Returns true if the first non-completed item's title is different fn _did_todo_item_change(&self, previous_todo: &Option) -> bool { - // Get first non-completed item from previous TODO list let prev_active = previous_todo.as_ref().and_then(|list| { - list.items.iter() + list.items + .iter() .find(|item| item.status != TodoListStatus::Completed) .map(|item| &item.title) }); // Get first non-completed item from current TODO list let curr_active = self.chain.todo_list.as_ref().and_then(|list| { - list.items.iter() + list.items + .iter() .find(|item| item.status != TodoListStatus::Completed) .map(|item| &item.title) }); @@ -396,4 +406,4 @@ impl Workflow { // If both exist but titles differ, change occurred prev_active != curr_active } -} \ No newline at end of file +} diff --git a/src/infrastructure/api_clients/openai/client.rs b/src/infrastructure/api_clients/openai/client.rs index d2eb83f..7c4b107 100644 --- a/src/infrastructure/api_clients/openai/client.rs +++ b/src/infrastructure/api_clients/openai/client.rs @@ -3,6 +3,15 @@ use super::translator::{build_llm_result, build_request_dto}; use crate::infrastructure::inference::LLMInferenceResult; use std::error::Error; +#[derive(Clone)] +pub enum AuthToken { + ApiKey(String), + OAuth { + token: String, + account_id: Option, + }, +} + #[derive(Debug)] pub enum OpenAIClientError { Api { @@ -16,6 +25,17 @@ pub enum OpenAIClientError { NoText { body: String, }, + BodyRead { + source: reqwest::Error, + status: reqwest::StatusCode, + content_encoding: Option, + content_type: Option, + content_length: Option, + transfer_encoding: Option, + response_headers: String, + partial_body: Option, + }, + AuthExpired, } impl std::fmt::Display for OpenAIClientError { @@ -34,6 +54,34 @@ impl std::fmt::Display for OpenAIClientError { OpenAIClientError::NoText { body } => { write!(f, "No output_text found in OpenAI response (body={})", body) } + OpenAIClientError::BodyRead { + source, + status, + content_encoding, + content_type, + content_length, + transfer_encoding, + response_headers, + partial_body, + } => { + write!( + f, + "Failed to read SSE response body (status={}, content-encoding={}, content-type={}, content-length={}, transfer-encoding={}, headers={}, partial-body={}, error={:?})", + status, + content_encoding.as_deref().unwrap_or(""), + content_type.as_deref().unwrap_or(""), + content_length + .map(|value| value.to_string()) + .unwrap_or_else(|| "".to_string()), + transfer_encoding.as_deref().unwrap_or(""), + response_headers, + partial_body.as_deref().unwrap_or(""), + source + ) + } + OpenAIClientError::AuthExpired => { + write!(f, "OAuth token expired - please re-authenticate") + } } } } @@ -42,23 +90,27 @@ impl std::error::Error for OpenAIClientError { fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { match self { OpenAIClientError::Deserialize { source, .. } => Some(source), + OpenAIClientError::BodyRead { source, .. } => Some(source), _ => None, } } } pub struct OpenAIClient { - api_key: String, + auth_token: AuthToken, model: String, client: reqwest::blocking::Client, } impl OpenAIClient { - pub fn new(api_key: String, model: String) -> Self { + pub fn new(auth_token: AuthToken, model: String) -> Self { + let client = reqwest::blocking::Client::builder() + .build() + .expect("Failed to build OpenAI HTTP client"); Self { - api_key, + auth_token, model, - client: reqwest::blocking::Client::new(), + client, } } @@ -71,12 +123,7 @@ impl OpenAIClient { ) -> Result> { let max_attempts = 3; for attempt in 1..=max_attempts { - match self.call_responses_api_inner( - tools, - chain, - images, - tracer.as_deref_mut(), - ) { + match self.call_responses_api_inner(tools, chain, images, tracer.as_deref_mut()) { Ok(result) => return Ok(result), Err(e) => { let error_str = e.to_string(); @@ -122,15 +169,14 @@ impl OpenAIClient { images: &[String], mut tracer: Option<&mut openai_agents_tracing::TracingFacade>, ) -> Result> { - let url = "https://api.openai.com/v1/responses"; + // Determine API URL based on auth type + let url = match &self.auth_token { + AuthToken::OAuth { .. } => "https://chatgpt.com/backend-api/codex/responses", + AuthToken::ApiKey(_) => "https://api.openai.com/v1/responses", + }; - let request_body = build_request_dto( - &self.model, - images, - tools, - chain, - tracer.as_deref_mut(), - ); + let request_body = + build_request_dto(&self.model, images, tools, chain, tracer.as_deref_mut()); // Start span with model name and add request as JSON if let Some(tracer) = &mut tracer { @@ -142,30 +188,97 @@ impl OpenAIClient { } } - let response = self - .client - .post(url) - .header("Authorization", format!("Bearer {}", self.api_key)) - .json(&request_body) + // Build request with appropriate headers + let mut request_builder = self.client.post(url).json(&request_body); + + // Add authentication header + match &self.auth_token { + AuthToken::ApiKey(key) => { + request_builder = + request_builder.header("Authorization", format!("Bearer {}", key)); + } + AuthToken::OAuth { token, account_id } => { + request_builder = + request_builder.header("Authorization", format!("Bearer {}", token)); + if let Some(acc_id) = account_id { + request_builder = request_builder.header("ChatGPT-Account-Id", acc_id); + } + } + } + + let mut response = request_builder + .header("Accept", "text/event-stream") + .header(reqwest::header::ACCEPT_ENCODING, "identity") .send()?; let status = response.status(); - let body = response.text()?; + let content_encoding = response + .headers() + .get(reqwest::header::CONTENT_ENCODING) + .and_then(|value| value.to_str().ok()) + .map(|value| value.to_string()); + let content_type = response + .headers() + .get(reqwest::header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .map(|value| value.to_string()); + let content_length = response.content_length(); + let transfer_encoding = response + .headers() + .get(reqwest::header::TRANSFER_ENCODING) + .and_then(|value| value.to_str().ok()) + .map(|value| value.to_string()); + let response_headers = Self::format_headers_excerpt(response.headers()); if !status.is_success() { + let body = response + .text() + .unwrap_or_else(|e| format!("(failed to read body: {})", e)); if let Some(t) = tracer { t.end_span(&self.model); } + + // Check for OAuth expiration + if status.as_u16() == 401 { + return Err(Box::new(OpenAIClientError::AuthExpired)); + } + return Err(Box::new(OpenAIClientError::Api { status, body })); } - let dto = match serde_json::from_str::(&body) { + // Read SSE stream as raw bytes to avoid encoding issues with chunked transfer + let mut body_bytes = Vec::new(); + if let Err(e) = response.copy_to(&mut body_bytes) { + let partial_body = if body_bytes.is_empty() { + None + } else { + Some(Self::format_body_excerpt( + &String::from_utf8_lossy(&body_bytes), + 4000, + )) + }; + return Err(Box::new(OpenAIClientError::BodyRead { + source: e, + status, + content_encoding: content_encoding.clone(), + content_type: content_type.clone(), + content_length, + transfer_encoding: transfer_encoding.clone(), + response_headers: response_headers.clone(), + partial_body, + })); + } + let body = String::from_utf8_lossy(&body_bytes).into_owned(); + log::debug!("SSE response length: {} bytes", body.len()); + log::trace!("SSE response body: {}", &body[..body.len().min(2000)]); + + let dto = match Self::parse_sse_response(&body) { Ok(v) => v, Err(e) => { if let Some(t) = tracer { t.end_span(&self.model); } - return Err(Box::new(OpenAIClientError::Deserialize { source: e, body })); + return Err(e); } }; @@ -179,9 +292,229 @@ impl OpenAIClient { let result = build_llm_result(dto, tools); if result.summary.is_empty() && result.tool_call.is_none() { - return Err(Box::new(OpenAIClientError::NoText { body })); + return Err(Box::new(OpenAIClientError::NoText { + body: format!( + "empty response from SSE stream; sse_body={}", + Self::format_body_excerpt(&body, 4000) + ), + })); } Ok(result) } + + /// Parse an SSE response body, extracting the ResponseDTO from the `response.completed` event. + /// + /// SSE format from OpenAI Responses API: + /// ```text + /// event: response.created + /// data: {"type":"response.created","response":{...}} + /// + /// event: response.output_text.delta + /// data: {"type":"response.output_text.delta","delta":"Hello",...} + /// + /// event: response.completed + /// data: {"type":"response.completed","response":{"output":[...],...}} + /// ``` + fn parse_sse_response(body: &str) -> Result> { + let mut current_event: Option = None; + let mut completed_data: Option = None; + + for line in body.lines() { + if let Some(event_name) = line.strip_prefix("event:") { + current_event = Some(event_name.trim().to_string()); + } else if let Some(data_payload) = line.strip_prefix("data:") { + let data = data_payload.trim_start(); + if current_event.as_deref() == Some("response.completed") { + completed_data = Some(data.to_string()); + break; + } + } + // Empty lines are event boundaries in SSE — reset current_event + if line.is_empty() { + current_event = None; + } + } + + let data = match completed_data { + Some(d) => d, + None => { + // Fallback: maybe the body is plain JSON (non-SSE), try parsing directly + log::warn!( + "No response.completed event found in SSE stream, trying direct JSON parse" + ); + return serde_json::from_str::(body.trim()).map_err(|e| { + log::error!( + "Direct JSON parse also failed. Body prefix: {}", + &body[..body.len().min(500)] + ); + Box::new(OpenAIClientError::Deserialize { + source: e, + body: format!("sse_body={}", Self::format_body_excerpt(body, 4000)), + }) as Box + }); + } + }; + + // The response.completed data has shape: {"type":"response.completed","response":{...}} + let wrapper: serde_json::Value = serde_json::from_str(&data).map_err(|e| { + log::error!( + "Failed to parse response.completed JSON: {}", + &data[..data.len().min(500)] + ); + Box::new(OpenAIClientError::Deserialize { + source: e, + body: format!( + "response.completed data={}; sse_body={}", + Self::format_body_excerpt(&data, 1000), + Self::format_body_excerpt(body, 4000) + ), + }) as Box + })?; + + // Extract the "response" field which contains the actual response object + let response_obj = wrapper.get("response").unwrap_or(&wrapper); + + serde_json::from_value::(response_obj.clone()).map_err(|e| { + log::error!( + "Failed to deserialize ResponseDTO from: {}", + &response_obj.to_string()[..response_obj.to_string().len().min(500)] + ); + Box::new(OpenAIClientError::Deserialize { + source: e, + body: format!( + "response_obj={}; sse_body={}", + Self::format_body_excerpt(&response_obj.to_string(), 1000), + Self::format_body_excerpt(body, 4000) + ), + }) as Box + }) + } + + fn format_body_excerpt(body: &str, max_len: usize) -> String { + let trimmed = body.trim(); + let total = trimmed.chars().count(); + if total <= max_len { + trimmed.to_string() + } else { + let excerpt: String = trimmed.chars().take(max_len).collect(); + format!( + "{}...[truncated {} chars]", + excerpt, + total.saturating_sub(max_len) + ) + } + } + + fn format_headers_excerpt(headers: &reqwest::header::HeaderMap) -> String { + let mut rendered = String::new(); + for (name, value) in headers.iter() { + if !rendered.is_empty() { + rendered.push_str(", "); + } + let value_str = value.to_str().unwrap_or(""); + rendered.push_str(name.as_str()); + rendered.push('='); + rendered.push_str(value_str); + if rendered.len() > 2000 { + rendered.push_str("...[truncated]"); + break; + } + } + if rendered.is_empty() { + "".to_string() + } else { + rendered + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parse_sse_text_response() { + let sse_body = "\ +event: response.created\n\ +data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_1\",\"output\":[]}}\n\ +\n\ +event: response.output_text.delta\n\ +data: {\"type\":\"response.output_text.delta\",\"delta\":\"Hello\"}\n\ +\n\ +event: response.completed\n\ +data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"output\":[{\"type\":\"message\",\"content\":[{\"type\":\"output_text\",\"text\":\"Hello, world!\"}]}]}}\n\ +\n"; + + let dto = OpenAIClient::parse_sse_response(sse_body).expect("should parse SSE"); + let (summary, tool_call) = dto.extract_parts(); + assert_eq!(summary, "Hello, world!"); + assert!(tool_call.is_none()); + } + + #[test] + fn parse_sse_function_call_response() { + let sse_body = "\ +event: response.created\n\ +data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_2\",\"output\":[]}}\n\ +\n\ +event: response.completed\n\ +data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_2\",\"output\":[{\"type\":\"function_call\",\"name\":\"shell_exec\",\"arguments\":\"{\\\"command\\\":\\\"ls\\\"}\",\"call_id\":\"call_abc\"}]}}\n\ +\n"; + + let dto = OpenAIClient::parse_sse_response(sse_body).expect("should parse SSE"); + let (summary, tool_call) = dto.extract_parts(); + assert!(summary.is_empty()); + let call = tool_call.expect("should have tool call"); + assert_eq!(call.name, "shell_exec"); + assert_eq!(call.call_id, "call_abc"); + } + + #[test] + fn parse_sse_mixed_text_and_function_call() { + let sse_body = "\ +event: response.completed\n\ +data: {\"type\":\"response.completed\",\"response\":{\"output\":[{\"type\":\"message\",\"content\":[{\"type\":\"output_text\",\"text\":\"Let me check.\"}]},{\"type\":\"function_call\",\"name\":\"read_file\",\"arguments\":\"{}\",\"call_id\":\"call_xyz\"}]}}\n\ +\n"; + + let dto = OpenAIClient::parse_sse_response(sse_body).expect("should parse SSE"); + let (summary, tool_call) = dto.extract_parts(); + assert_eq!(summary, "Let me check."); + assert!(tool_call.is_some()); + } + + #[test] + fn parse_sse_no_completed_event_falls_back_to_direct_json() { + // If body is plain JSON (not SSE), should still work + let json_body = "{\"output\":[{\"type\":\"message\",\"content\":[{\"type\":\"output_text\",\"text\":\"Direct JSON\"}]}]}"; + + let dto = OpenAIClient::parse_sse_response(json_body).expect("should parse direct JSON"); + let (summary, _) = dto.extract_parts(); + assert_eq!(summary, "Direct JSON"); + } + + #[test] + fn parse_sse_empty_body_errors() { + let result = OpenAIClient::parse_sse_response(""); + assert!(result.is_err()); + } + + #[test] + fn parse_sse_ignores_earlier_events() { + // Ensure we pick up data from response.completed, not from earlier events + let sse_body = "\ +event: response.created\n\ +data: {\"type\":\"response.created\",\"response\":{\"id\":\"r1\",\"output\":[]}}\n\ +\n\ +event: response.in_progress\n\ +data: {\"type\":\"response.in_progress\"}\n\ +\n\ +event: response.completed\n\ +data: {\"type\":\"response.completed\",\"response\":{\"output\":[{\"type\":\"message\",\"content\":[{\"type\":\"output_text\",\"text\":\"Final answer\"}]}]}}\n\ +\n"; + + let dto = OpenAIClient::parse_sse_response(sse_body).expect("should parse"); + let (summary, _) = dto.extract_parts(); + assert_eq!(summary, "Final answer"); + } } diff --git a/src/infrastructure/api_clients/openai/dto/request.rs b/src/infrastructure/api_clients/openai/dto/request.rs index a747bc4..dff01a1 100644 --- a/src/infrastructure/api_clients/openai/dto/request.rs +++ b/src/infrastructure/api_clients/openai/dto/request.rs @@ -141,7 +141,7 @@ impl RequestDTO { tool_choice: "auto".to_string(), parallel_tool_calls: false, store: false, - stream: false, + stream: true, } } } diff --git a/src/infrastructure/auth/mod.rs b/src/infrastructure/auth/mod.rs new file mode 100644 index 0000000..609c0dc --- /dev/null +++ b/src/infrastructure/auth/mod.rs @@ -0,0 +1,2 @@ +pub mod oauth; +pub use oauth::{initiate_oauth_flow, refresh_oauth_tokens}; \ No newline at end of file diff --git a/src/infrastructure/auth/oauth.rs b/src/infrastructure/auth/oauth.rs new file mode 100644 index 0000000..faa5435 --- /dev/null +++ b/src/infrastructure/auth/oauth.rs @@ -0,0 +1,224 @@ +use oauth2::{ + AuthUrl, AuthorizationCode, ClientId, CsrfToken, PkceCodeChallenge, + PkceCodeVerifier, RedirectUrl, RefreshToken, Scope, TokenResponse, TokenUrl, +}; +use oauth2::basic::BasicClient; +use oauth2::reqwest::http_client; +use sha2::{Digest, Sha256}; +use std::sync::mpsc; +use std::thread; +use std::time::Duration; +use tiny_http::{Response, Server}; + +const CLIENT_ID: &str = "app_EMoamEEZ73f0CkXaXp7hrann"; +const AUTH_URL: &str = "https://auth.openai.com/oauth/authorize"; +const TOKEN_URL: &str = "https://auth.openai.com/oauth/token"; +const REDIRECT_URI: &str = "http://localhost:1455/auth/callback"; +const SCOPES: &[&str] = &["openid", "profile", "email", "offline_access"]; + +pub struct OAuthTokens { + pub access_token: String, + pub refresh_token: String, + pub expires_in: i64, + pub account_id: Option, +} + +/// Generates a cryptographically secure random string for PKCE +fn generate_random_string(length: usize) -> String { + use rand::Rng; + const CHARSET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-._~"; + let mut rng = rand::thread_rng(); + (0..length) + .map(|_| CHARSET[rng.gen_range(0..CHARSET.len())] as char) + .collect() +} + +/// Initiates OAuth flow and returns tokens +pub fn initiate_oauth_flow() -> Result { + // Generate PKCE code verifier and challenge + let code_verifier = generate_random_string(43); + let mut hasher = Sha256::new(); + hasher.update(code_verifier.as_bytes()); + + use base64::Engine; + let _challenge = base64::engine::general_purpose::URL_SAFE_NO_PAD + .encode(hasher.finalize()); + + let pkce_verifier = PkceCodeVerifier::new(code_verifier); + let pkce_challenge = PkceCodeChallenge::from_code_verifier_sha256(&pkce_verifier); + + // Build OAuth client + let client = BasicClient::new( + ClientId::new(CLIENT_ID.to_string()), + None, + AuthUrl::new(AUTH_URL.to_string()).map_err(|e| e.to_string())?, + Some(TokenUrl::new(TOKEN_URL.to_string()).map_err(|e| e.to_string())?), + ) + .set_redirect_uri(RedirectUrl::new(REDIRECT_URI.to_string()).map_err(|e| e.to_string())?); + + // Generate authorization URL with OpenCode-specific parameters + let (mut auth_url, csrf_state) = client + .authorize_url(CsrfToken::new_random) + .add_scopes(SCOPES.iter().map(|s| Scope::new(s.to_string()))) + .set_pkce_challenge(pkce_challenge) + .url(); + + // Add OpenCode-specific query parameters + auth_url.query_pairs_mut() + .append_pair("id_token_add_organizations", "true") + .append_pair("codex_cli_simplified_flow", "true") + .append_pair("originator", "drastis"); + + // Start local HTTP server for callback + let server = Server::http("127.0.0.1:1455") + .map_err(|e| format!("Failed to start callback server on port 1455: {}", e))?; + + log::info!("Starting OAuth flow, opening browser..."); + + // Open browser + open::that(auth_url.to_string()) + .map_err(|e| format!("Failed to open browser: {}", e))?; + + // Wait for callback with timeout + let (tx, rx) = mpsc::channel(); + let csrf_secret = csrf_state.secret().clone(); + + thread::spawn(move || { + // Listen for callback (2 minute timeout) + if let Ok(Some(request)) = server.recv_timeout(Duration::from_secs(120)) { + let url_str = format!("http://localhost{}", request.url()); + + // Send success HTML response + let html = r#" + + Authentication Successful + +

✓ Authentication Successful

+

You can close this window and return to the CLI.

+ + + + "#; + let _ = request.respond(Response::from_string(html) + .with_header(tiny_http::Header::from_bytes(&b"Content-Type"[..], &b"text/html"[..]).unwrap())); + + let _ = tx.send(url_str); + } + }); + + // Wait for callback + let callback_url = rx.recv_timeout(Duration::from_secs(120)) + .map_err(|_| "OAuth timeout: No response received within 2 minutes".to_string())?; + + // Parse callback URL + let url = url::Url::parse(&callback_url) + .map_err(|e| format!("Invalid callback URL: {}", e))?; + + // Extract code and state from query parameters + let mut code = None; + let mut state = None; + + for (key, value) in url.query_pairs() { + match key.as_ref() { + "code" => code = Some(value.to_string()), + "state" => state = Some(value.to_string()), + _ => {} + } + } + + let code = code.ok_or("No authorization code in callback")?; + let state = state.ok_or("No state in callback")?; + + // Verify CSRF token + if state != csrf_secret { + return Err("CSRF token mismatch - possible security attack".to_string()); + } + + // Exchange code for tokens + log::info!("Exchanging authorization code for tokens..."); + + let token_response = client + .exchange_code(AuthorizationCode::new(code)) + .set_pkce_verifier(pkce_verifier) + .request(http_client) + .map_err(|e| format!("Token exchange failed: {}", e))?; + + let access_token = token_response.access_token().secret().clone(); + let refresh_token = token_response.refresh_token() + .ok_or("No refresh token received")? + .secret() + .clone(); + let expires_in = token_response.expires_in() + .ok_or("No expiration time received")? + .as_secs() as i64; + + // Extract account ID from access token (JWT) + let account_id = extract_account_id_from_token(&access_token); + + log::info!("OAuth authentication successful!"); + + Ok(OAuthTokens { + access_token, + refresh_token, + expires_in, + account_id, + }) +} + +/// Refreshes OAuth tokens using refresh token +pub fn refresh_oauth_tokens(refresh_token: &str) -> Result { + let client = BasicClient::new( + ClientId::new(CLIENT_ID.to_string()), + None, + AuthUrl::new(AUTH_URL.to_string()).map_err(|e| e.to_string())?, + Some(TokenUrl::new(TOKEN_URL.to_string()).map_err(|e| e.to_string())?), + ); + + log::info!("Refreshing OAuth tokens..."); + + let token_response = client + .exchange_refresh_token(&RefreshToken::new(refresh_token.to_string())) + .request(http_client) + .map_err(|e| format!("Token refresh failed: {}", e))?; + + let access_token = token_response.access_token().secret().clone(); + let refresh_token = token_response.refresh_token() + .map(|t| t.secret().clone()) + .unwrap_or_else(|| refresh_token.to_string()); + let expires_in = token_response.expires_in() + .ok_or("No expiration time received")? + .as_secs() as i64; + + let account_id = extract_account_id_from_token(&access_token); + + Ok(OAuthTokens { + access_token, + refresh_token, + expires_in, + account_id, + }) +} + +/// Extracts ChatGPT account ID from JWT token +fn extract_account_id_from_token(token: &str) -> Option { + use base64::Engine; + + // JWT format: header.payload.signature + let parts: Vec<&str> = token.split('.').collect(); + if parts.len() != 3 { + return None; + } + + // Decode payload (base64url) + let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(parts[1]) + .ok()?; + + // Parse JSON + let json: serde_json::Value = serde_json::from_slice(&payload).ok()?; + + // Extract chatgpt_account_id + json.get("chatgpt_account_id") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()) +} diff --git a/src/infrastructure/cli/actions/mod.rs b/src/infrastructure/cli/actions/mod.rs index b8bd6b4..5f30075 100644 --- a/src/infrastructure/cli/actions/mod.rs +++ b/src/infrastructure/cli/actions/mod.rs @@ -1,3 +1,4 @@ pub mod change_settings; pub mod insert_openai_key; +pub mod oauth_login; pub mod select_openai_model; diff --git a/src/infrastructure/cli/actions/oauth_login.rs b/src/infrastructure/cli/actions/oauth_login.rs new file mode 100644 index 0000000..72a8adf --- /dev/null +++ b/src/infrastructure/cli/actions/oauth_login.rs @@ -0,0 +1,48 @@ +use crate::infrastructure::auth::initiate_oauth_flow; +use crate::infrastructure::cli::repl::refresh_settings_from_db; +use crate::infrastructure::cli::state::UiState; +use crate::infrastructure::db::DbPool; +use crate::infrastructure::event_bus::{EventBus, UiToAgentEvent}; +use crate::repository::UserSettingsRepository; + +pub fn start_oauth_flow( + bus: &EventBus, + conn: &DbPool, + state: &mut UiState, +) -> Result<(), String> { + log::info!("Starting OpenAI OAuth login flow..."); + + // Initiate OAuth flow (blocks until complete or timeout) + let tokens = initiate_oauth_flow()?; + + // Store tokens in database + let conn_guard = conn + .get() + .map_err(|e| format!("Database connection failed: {}", e))?; + let settings_repo = UserSettingsRepository::new(&*conn_guard); + + settings_repo.update_oauth_tokens( + &tokens.access_token, + &tokens.refresh_token, + tokens.expires_in, + tokens.account_id.as_deref(), + )?; + + log::info!("OAuth tokens stored successfully"); + + // Notify event controller to reload engine with OAuth credentials + let _ = bus.ui_to_agent_tx.send(UiToAgentEvent::SettingsUpdateEvent { + model: None, + openai_api_key: None, + use_behavior_trees: None, + openai_tracing_enabled: None, + web_search_enabled: None, + max_tool_calls_per_request: None, + brave_api_key: None, + }); + + // Refresh UI state from database + refresh_settings_from_db(conn, state)?; + + Ok(()) +} diff --git a/src/infrastructure/cli/actions/select_openai_model.rs b/src/infrastructure/cli/actions/select_openai_model.rs index f57fd94..536de31 100644 --- a/src/infrastructure/cli/actions/select_openai_model.rs +++ b/src/infrastructure/cli/actions/select_openai_model.rs @@ -35,7 +35,7 @@ pub fn handle_model_entry( refresh_settings_from_db(conn, state)?; state.mode = UiMode::Normal; } - ModelEntry::OpenAi { name } => { + ModelEntry::OpenAi { name, .. } => { update_model_selection_in_db(conn, &ModelSelection::OpenAiModel(name.clone()))?; let _ = bus .ui_to_agent_tx @@ -52,17 +52,7 @@ pub fn handle_model_entry( state.mode = UiMode::Normal; } ModelEntry::OpenAiSelect => { - if insert_openai_key::api_key_missing(conn)? { - insert_openai_key::begin_prompt(state, false); - } else { - state.models.openai_available_status = LoadStatus::Loading; - state.openai_fetch_pending = true; - state.mode = UiMode::Popup(PopupState::OpenAiAvailable { - selected: 0, - filter: String::new(), - cursor: 0, - }); - } + state.mode = UiMode::Popup(PopupState::AuthMethodSelect { selected: 0 }); } ModelEntry::Loading | ModelEntry::Empty => {} } @@ -70,6 +60,42 @@ pub fn handle_model_entry( Ok(()) } +pub fn handle_auth_method_api_key( + conn: &DbPool, + state: &mut UiState, +) -> Result<(), String> { + if insert_openai_key::api_key_missing(conn)? { + insert_openai_key::begin_prompt(state, false); + } else { + state.models.openai_available_status = LoadStatus::Loading; + state.openai_fetch_pending = true; + state.mode = UiMode::Popup(PopupState::OpenAiAvailable { + selected: 0, + filter: String::new(), + cursor: 0, + }); + } + Ok(()) +} + +pub fn handle_auth_method_oauth( + bus: &EventBus, + conn: &DbPool, + state: &mut UiState, +) -> Result<(), String> { + use crate::infrastructure::cli::actions::oauth_login; + oauth_login::start_oauth_flow(bus, conn, state)?; + // After successful OAuth, proceed to model selection + state.models.openai_available_status = LoadStatus::Loading; + state.openai_fetch_pending = true; + state.mode = UiMode::Popup(PopupState::OpenAiAvailable { + selected: 0, + filter: String::new(), + cursor: 0, + }); + Ok(()) +} + pub fn handle_openai_available_selection( bus: &EventBus, conn: &DbPool, @@ -91,4 +117,4 @@ pub fn handle_openai_available_selection( refresh_settings_from_db(conn, state)?; state.mode = UiMode::Normal; Ok(()) -} +} \ No newline at end of file diff --git a/src/infrastructure/cli/components/permissions.rs b/src/infrastructure/cli/components/permissions.rs index 7a213e9..3b3ecc2 100644 --- a/src/infrastructure/cli/components/permissions.rs +++ b/src/infrastructure/cli/components/permissions.rs @@ -16,7 +16,6 @@ pub fn render( scope: &str, selected: usize, theme: Theme, - is_read_only: bool, ) { let width = min( (size.width as f32 * 0.7) as u16, @@ -58,7 +57,7 @@ pub fn render( width: area.width, height: 5, }; - let options = permission_options(command.is_some(), is_read_only); + let options = permission_options(command.is_some()); let items: Vec = options .iter() .map(|option| ListItem::new(Line::from(Span::raw(*option)))) @@ -75,23 +74,17 @@ pub fn render( frame.render_stateful_widget(list, options_area, &mut list_state); } -fn permission_options(has_command: bool, is_read_only: bool) -> Vec<&'static str> { +fn permission_options(has_command: bool) -> Vec<&'static str> { if has_command { vec![ "[A] Allow", "[C] Allow this Command for this Project", "[D] Disallow", ] - } else if is_read_only { - vec![ - "[A] Allow", - "[R] Allow All Reads in Project for this Session", - "[D] Disallow", - ] } else { vec![ "[A] Allow", - "[W] Allow All Writes in Project for this Session", + "[W] Allow All Writes in this Project", "[D] Disallow", ] } diff --git a/src/infrastructure/cli/components/settings_panel.rs b/src/infrastructure/cli/components/settings_panel.rs index 78f06f2..4d0ac1c 100644 --- a/src/infrastructure/cli/components/settings_panel.rs +++ b/src/infrastructure/cli/components/settings_panel.rs @@ -1,3 +1,4 @@ +use crate::domain::AuthMethod; use crate::infrastructure::cli::helpers::{checkbox_item, list_state, panel_block}; use crate::infrastructure::cli::theme::{Theme, PANEL_PADDING}; use ratatui::style::{Modifier, Style}; @@ -13,8 +14,29 @@ pub fn render( openai_tracing: bool, web_search: bool, max_tool_calls: i32, + auth_method: &AuthMethod, + oauth_token_expiry: Option, theme: Theme, ) { + // Determine auth status display + let auth_status = match auth_method { + AuthMethod::OAuth => { + let expiry_msg = if let Some(exp) = oauth_token_expiry { + let now = chrono::Utc::now().timestamp(); + if exp > now { + let mins_left = (exp - now) / 60; + format!(" (expires in {}m)", mins_left) + } else { + " (EXPIRED - re-login needed)".to_string() + } + } else { + "".to_string() + }; + format!("OAuth{}", expiry_msg) + } + AuthMethod::ApiKey => "API Key".to_string(), + }; + let items = vec![ checkbox_item("Behavior trees", behavior_trees), checkbox_item("OpenAI tracing", openai_tracing), @@ -26,6 +48,13 @@ pub fn render( Style::default().fg(theme.active), ), ])), + ListItem::new(Line::from(vec![ + Span::raw("Auth: "), + Span::styled( + auth_status, + Style::default().fg(theme.active), + ), + ])), ]; let list = List::new(items) .highlight_style( diff --git a/src/infrastructure/cli/controls.rs b/src/infrastructure/cli/controls.rs index a4b41cc..5d87156 100644 --- a/src/infrastructure/cli/controls.rs +++ b/src/infrastructure/cli/controls.rs @@ -1,4 +1,4 @@ -use crate::domain::permissions::PermissionDecision; +use crate::domain::permissions::UserPermissionDecision; use crate::domain::AgentModeType; use crate::infrastructure::cli::actions::{ change_settings, insert_openai_key, select_openai_model, @@ -24,6 +24,13 @@ pub fn handle_key_event( state: &mut UiState, key: KeyEvent, ) -> Result<(), String> { + if key.code == KeyCode::Char('c') && key.modifiers.contains(KeyModifiers::CONTROL) { + if state.request_in_flight.is_some() { + let _ = bus.ui_to_agent_tx.send(UiToAgentEvent::CancelRequest); + } + return Ok(()); + } + let mode_snapshot = state.mode.clone(); match mode_snapshot { UiMode::Normal => handle_normal_key(bus, conn, state, key), @@ -213,8 +220,8 @@ fn handle_normal_key( // Enter right_container review mode for whatever is displayed let has_file_changes = state.file_changes.is_some() && !state.file_changes.as_ref().unwrap().changes.is_empty(); - let has_todo_list = state.todo_list.is_some() - && !state.todo_list.as_ref().unwrap().items.is_empty(); + let has_todo_list = + state.todo_list.is_some() && !state.todo_list.as_ref().unwrap().items.is_empty(); // Priority: file changes first (if both exist), otherwise whichever exists if has_file_changes { @@ -237,8 +244,8 @@ fn handle_normal_key( // Enter right_container review mode for whatever is displayed let has_file_changes = state.file_changes.is_some() && !state.file_changes.as_ref().unwrap().changes.is_empty(); - let has_todo_list = state.todo_list.is_some() - && !state.todo_list.as_ref().unwrap().items.is_empty(); + let has_todo_list = + state.todo_list.is_some() && !state.todo_list.as_ref().unwrap().items.is_empty(); // Priority: file changes first (if both exist), otherwise whichever exists if has_file_changes { @@ -568,12 +575,17 @@ fn handle_popup_key( } _ => {} }, - PopupState::PermissionPrompt { selected, is_read_only, command, .. } => { + PopupState::PermissionPrompt { + selected, + is_read_only, + command, + .. + } => { let options_len = permissions::option_count(); let has_command = command.is_some(); match key.code { KeyCode::Esc | KeyCode::Char('d') | KeyCode::Char('D') => { - submit_permission_decision(bus, popup, PermissionDecision::Ask)?; + submit_permission_decision(bus, popup, UserPermissionDecision::Deny)?; state.mode = UiMode::Normal; } KeyCode::Up => { @@ -587,39 +599,45 @@ fn handle_popup_key( } } KeyCode::Char('a') | KeyCode::Char('A') => { - submit_permission_decision(bus, popup, PermissionDecision::AllowOnce)?; + submit_permission_decision(bus, popup, UserPermissionDecision::AllowOnce)?; state.mode = UiMode::Normal; } KeyCode::Char('c') | KeyCode::Char('C') => { if has_command { - submit_permission_decision(bus, popup, PermissionDecision::AllowCommandForProject)?; + submit_permission_decision( + bus, + popup, + UserPermissionDecision::AlwaysAllow, + )?; state.mode = UiMode::Normal; } } KeyCode::Char('r') | KeyCode::Char('R') => { if !has_command && *is_read_only { - submit_permission_decision(bus, popup, PermissionDecision::AllowAllReadsInSession)?; + submit_permission_decision( + bus, + popup, + UserPermissionDecision::AlwaysAllow, + )?; state.mode = UiMode::Normal; } } KeyCode::Char('w') | KeyCode::Char('W') => { if !has_command && !*is_read_only { - submit_permission_decision(bus, popup, PermissionDecision::AllowAllWritesInSession)?; + submit_permission_decision( + bus, + popup, + UserPermissionDecision::AlwaysAllow, + )?; state.mode = UiMode::Normal; } } KeyCode::Enter => { let decision = match *selected { - 0 => PermissionDecision::AllowOnce, - 1 => if has_command { - PermissionDecision::AllowCommandForProject - } else if *is_read_only { - PermissionDecision::AllowAllReadsInSession - } else { - PermissionDecision::AllowAllWritesInSession - }, - 2 => PermissionDecision::Ask, - _ => PermissionDecision::Ask, + 0 => UserPermissionDecision::AllowOnce, + 1 => UserPermissionDecision::AlwaysAllow, + 2 => UserPermissionDecision::Deny, + _ => UserPermissionDecision::Deny, }; submit_permission_decision(bus, popup, decision)?; state.mode = UiMode::Normal; @@ -627,6 +645,29 @@ fn handle_popup_key( _ => {} } } + PopupState::AuthMethodSelect { selected } => match key.code { + KeyCode::Esc => state.mode = UiMode::Popup(PopupState::ModelSelect { selected: 0 }), + KeyCode::Up => { + if *selected > 0 { + *selected -= 1; + } + } + KeyCode::Down => { + if *selected < 1 { + *selected += 1; + } + } + KeyCode::Enter => match *selected { + 0 => { + select_openai_model::handle_auth_method_api_key(conn, state)?; + } + 1 => { + select_openai_model::handle_auth_method_oauth(bus, conn, state)?; + } + _ => {} + }, + _ => {} + }, PopupState::ContinueSelect { sessions, selected } => { let count = sessions.len(); match key.code { @@ -644,9 +685,9 @@ fn handle_popup_key( KeyCode::Enter => { let session_id = sessions[*selected].id; state.session_id = Some(session_id); - let _ = bus.ui_to_agent_tx.send(UiToAgentEvent::SessionContinueEvent { - session_id, - }); + let _ = bus + .ui_to_agent_tx + .send(UiToAgentEvent::SessionContinueEvent { session_id }); state.mode = UiMode::Normal; } _ => {} @@ -862,7 +903,7 @@ fn handle_popup_input(state: &mut UiState, key: KeyEvent) { fn submit_permission_decision( bus: &EventBus, popup: &PopupState, - decision: PermissionDecision, + decision: UserPermissionDecision, ) -> Result<(), String> { if let PopupState::PermissionPrompt { request_id, .. } = popup { bus.ui_to_agent_tx @@ -1328,14 +1369,12 @@ fn open_continue_popup(conn: &DbPool, state: &mut UiState) { .and_then(|p| p.file_name().map(|n| n.to_string_lossy().to_string())) .unwrap_or_else(|| "default".to_string()); - let projects_repo = - crate::repository::ProjectsRepository::new(&*conn_guard); + let projects_repo = crate::repository::ProjectsRepository::new(&*conn_guard); let project = projects_repo .find_by_name(&project_name)? .ok_or_else(|| "not_found".to_string())?; - let sessions_repo = - crate::repository::SessionsRepository::new(&*conn_guard); + let sessions_repo = crate::repository::SessionsRepository::new(&*conn_guard); let rows = sessions_repo.find_by_project_recent(project.id, 5)?; if rows.is_empty() { diff --git a/src/infrastructure/cli/repl.rs b/src/infrastructure/cli/repl.rs index d17b00e..eec0886 100644 --- a/src/infrastructure/cli/repl.rs +++ b/src/infrastructure/cli/repl.rs @@ -1,4 +1,4 @@ -use crate::domain::{AgentModeType, ModelType}; +use crate::domain::{AgentModeType, ModelAuthType, ModelType}; use crate::infrastructure::cli::state::{ FileChangesDisplay, LoadStatus, PopupState, ProgressEntry, ProgressKind, RequestIndicator, RequestStatusDisplay, UiMode, UiState, @@ -180,9 +180,7 @@ fn handle_agent_event( let can_update = state .last_progress_update .map(|last| { - std::time::Instant::now() - .duration_since(last) - .as_millis() + std::time::Instant::now().duration_since(last).as_millis() >= crate::infrastructure::cli::state::MIN_PROGRESS_DISPLAY_MS }) .unwrap_or(true); // Update immediately if no previous update @@ -236,15 +234,29 @@ fn handle_agent_event( } } - // Only show summary for failures, skip for success - if let Some(summary) = summary { + // Show failures via summary; for successful requests, fallback to summary + // when final_message is missing/empty (common when resuming older sessions). + if let Some(summary_text) = summary { if status != RequestStatus::Success { state.push_progress(ProgressEntry { - text: summary, + text: summary_text, kind: finish_kind, active: false, request_id: None, }); + } else { + let final_is_empty = final_message + .as_ref() + .map(|m| m.trim().is_empty()) + .unwrap_or(true); + if final_is_empty { + state.push_progress(ProgressEntry { + text: summary_text, + kind: ProgressKind::Info, + active: false, + request_id: None, + }); + } } } if let Some(message) = final_message { @@ -339,6 +351,7 @@ pub(crate) fn refresh_models_from_db(conn: &DbPool, state: &mut UiState) -> Resu .map(|name| crate::infrastructure::event_bus::OpenAiModelInfo { _id: model.id, name, + auth_type: model.auth_type.as_str().to_string(), }) }) .collect(); @@ -362,6 +375,8 @@ pub(crate) fn refresh_settings_from_db(conn: &DbPool, state: &mut UiState) -> Re state.settings.openai_tracing_enabled = settings.openai_tracing_enabled; state.settings.web_search_enabled = settings.web_search_enabled; state.settings.max_tool_calls_per_request = settings.max_tool_calls_per_request; + state.settings.auth_method = crate::domain::AuthMethod::from_str(&settings.auth_method); + state.settings.oauth_token_expiry = settings.oauth_token_expiry; state.settings.status = LoadStatus::Loaded; // Get both model name and type @@ -437,10 +452,23 @@ pub(crate) fn update_model_selection_in_db( }) }); let model_id = match existing_model { - Some(model) => model.id, + Some(model) => { + if model.auth_type != ModelAuthType::Local { + models_repo + .update_auth_type(model.id, ModelAuthType::Local) + .map_err(|e| e.to_string())?; + } + model.id + } None => { models_repo - .create(ModelType::Local, None, Some(&gguf_path), None) + .create( + ModelType::Local, + None, + Some(&gguf_path), + None, + ModelAuthType::Local, + ) .map_err(|e| e.to_string())? .id } @@ -457,14 +485,24 @@ pub(crate) fn update_model_selection_in_db( let existing_models = models_repo .find_by_type(ModelType::OpenAI) .map_err(|e| e.to_string())?; + let settings = settings_repo.get_current().map_err(|e| e.to_string())?; + let auth_method = crate::domain::AuthMethod::from_str(&settings.auth_method); + let auth_type = ModelAuthType::from_auth_method(&auth_method); let model_id = match existing_models .iter() .find(|model| model.model_name.as_deref() == Some(model_name)) { - Some(model) => model.id, + Some(model) => { + if model.auth_type != auth_type { + models_repo + .update_auth_type(model.id, auth_type) + .map_err(|e| e.to_string())?; + } + model.id + } None => { models_repo - .create(ModelType::OpenAI, None, None, Some(model_name)) + .create(ModelType::OpenAI, None, None, Some(model_name), auth_type) .map_err(|e| e.to_string())? .id } @@ -491,18 +529,63 @@ pub(crate) fn fetch_openai_available_models( .map_err(|e| format!("Failed to get connection: {}", e))?; let settings_repo = UserSettingsRepository::new(&*conn_guard); let settings = settings_repo.get_current().map_err(|e| e.to_string())?; - let api_key = match settings.openai_api_key { - Some(key) => key, - None => { - state.models.openai_available = - vec!["OpenAI API key is required. Please set it in settings.".to_string()]; - state.models.openai_available_status = LoadStatus::Loaded; - return Ok(()); + + let auth_method = crate::domain::AuthMethod::from_str(&settings.auth_method); + let auth_token = if auth_method == crate::domain::AuthMethod::OAuth { + match (&settings.oauth_access_token, &settings.oauth_account_id) { + (Some(token), account_id) => { + crate::infrastructure::api_clients::openai::client::AuthToken::OAuth { + token: token.clone(), + account_id: account_id.clone(), + } + } + _ => { + state.models.openai_available = + vec!["OAuth token not found. Please re-authenticate.".to_string()]; + state.models.openai_available_status = LoadStatus::Loaded; + return Ok(()); + } + } + } else { + match &settings.openai_api_key { + Some(key) if !key.trim().is_empty() => { + crate::infrastructure::api_clients::openai::client::AuthToken::ApiKey(key.clone()) + } + _ => { + state.models.openai_available = + vec!["OpenAI API key is required. Please set it in settings.".to_string()]; + state.models.openai_available_status = LoadStatus::Loaded; + return Ok(()); + } } }; + // OAuth tokens can't call /v1/models — return known Codex/ChatGPT models instead + if auth_method == crate::domain::AuthMethod::OAuth { + state.models.openai_available = vec![ + "gpt-5.3-codex-spark".to_string(), + "gpt-5.3-codex".to_string(), + "gpt-5.2-codex".to_string(), + "gpt-5.2".to_string(), + "gpt-5.1-codex-max".to_string(), + "gpt-5.1-codex".to_string(), + "gpt-5.1".to_string(), + "gpt-5-codex".to_string(), + "gpt-5-codex-mini".to_string(), + "gpt-5".to_string(), + "gpt-4.1".to_string(), + "gpt-4.1-mini".to_string(), + "gpt-4o".to_string(), + "gpt-4o-mini".to_string(), + "o3".to_string(), + "o4-mini".to_string(), + ]; + state.models.openai_available_status = LoadStatus::Loaded; + return Ok(()); + } + state.models.openai_available_status = LoadStatus::Loading; - match OpenAIEngine::new_general(&api_key, "gpt-4") { + match OpenAIEngine::new_general(auth_token, "gpt-4") { Ok(openai) => match openai.fetch_available_models() { Ok(models) => { state.models.openai_available = models; @@ -637,6 +720,7 @@ mod tests { state.models.openai = vec![OpenAiModelInfo { _id: 1, name: "gpt-4".to_string(), + auth_type: "api_key".to_string(), }]; let entries = model_entries(&state); assert_eq!(entries.len(), 3); @@ -734,4 +818,4 @@ mod tests { assert_eq!(model.gguf_file_path.as_deref(), Some(canonical.as_str())); }); } -} +} \ No newline at end of file diff --git a/src/infrastructure/cli/state.rs b/src/infrastructure/cli/state.rs index 7110c62..f6c8ba9 100644 --- a/src/infrastructure/cli/state.rs +++ b/src/infrastructure/cli/state.rs @@ -1,4 +1,5 @@ use crate::domain::tools::FileChange; +use crate::domain::AuthMethod; use crate::domain::AgentModeType; use crate::infrastructure::cli::loading_bar::LoadingBar; use crate::infrastructure::event_bus::{LocalModelInfo, OpenAiModelInfo}; @@ -118,6 +119,8 @@ pub struct SettingsCache { pub openai_tracing_enabled: bool, pub web_search_enabled: bool, pub max_tool_calls_per_request: i32, + pub auth_method: AuthMethod, + pub oauth_token_expiry: Option, } #[derive(Debug, Clone)] @@ -206,6 +209,9 @@ pub enum PopupState { sessions: Vec, selected: usize, }, + AuthMethodSelect { + selected: usize, + }, } #[derive(Debug, Clone)] @@ -297,6 +303,8 @@ impl UiState { openai_tracing_enabled: false, web_search_enabled: false, max_tool_calls_per_request: 50, + auth_method: AuthMethod::ApiKey, + oauth_token_expiry: None, }, loading_bar: LoadingBar::new(), openai_fetch_pending: false, diff --git a/src/infrastructure/cli/views/main_view.rs b/src/infrastructure/cli/views/main_view.rs index e2bb3cd..8113658 100644 --- a/src/infrastructure/cli/views/main_view.rs +++ b/src/infrastructure/cli/views/main_view.rs @@ -144,7 +144,13 @@ fn render_popup( .map(|entry| { let label = match entry { ModelEntry::Local { name, .. } => format!("Local: {}", name), - ModelEntry::OpenAi { name } => format!("OpenAI: {}", name), + ModelEntry::OpenAi { name, auth_type } => { + if auth_type.trim().is_empty() { + format!("OpenAI: {}", name) + } else { + format!("OpenAI: {} ({})", name, auth_type) + } + } ModelEntry::OpenAiSelect => "OpenAI: Select New Model".to_string(), ModelEntry::Loading => "Loading models...".to_string(), ModelEntry::Empty => "No models available".to_string(), @@ -272,6 +278,8 @@ fn render_popup( *openai_tracing, *web_search, *max_tool_calls, + &state.settings.auth_method, + state.settings.oauth_token_expiry, theme, ); } @@ -423,7 +431,6 @@ fn render_popup( paths, scope, selected, - is_read_only, .. } => { permissions::render( @@ -435,9 +442,36 @@ fn render_popup( scope, *selected, theme, - *is_read_only, ); } + PopupState::AuthMethodSelect { selected } => { + let height = 7u16; + let width = min( + (size.width as f32 * 0.7) as u16, + size.width.saturating_sub(2), + ); + let area = centered_rect(size, width, height); + frame.render_widget(Clear, area); + + let api_key_item = ListItem::new(Line::from(Span::styled( + "API Key", + Style::default().fg(theme.info_text), + ))); + let oauth_item = ListItem::new(Line::from(Span::styled( + "ChatGPT Plus (OAuth)", + Style::default().fg(theme.info_text), + ))); + let list = List::new(vec![api_key_item, oauth_item]) + .highlight_style( + Style::default() + .fg(theme.active) + .add_modifier(Modifier::BOLD), + ) + .block(panel_block(theme, theme.panel, PANEL_PADDING)) + .highlight_symbol("> "); + let mut list_state = list_state(*selected); + frame.render_stateful_widget(list, area, &mut list_state); + } PopupState::ContinueSelect { sessions, selected } => { if sessions.is_empty() { return; @@ -496,7 +530,7 @@ fn render_popup( #[derive(Debug, Clone)] pub enum ModelEntry { Local { name: String, path: String }, - OpenAi { name: String }, + OpenAi { name: String, auth_type: String }, OpenAiSelect, Loading, Empty, @@ -517,6 +551,7 @@ pub fn model_entries(state: &UiState) -> Vec { for model in &state.models.openai { entries.push(ModelEntry::OpenAi { name: model.name.clone(), + auth_type: model.auth_type.clone(), }); } entries.push(ModelEntry::OpenAiSelect); diff --git a/src/infrastructure/db.rs b/src/infrastructure/db.rs index 2da4f27..bc6a9d5 100644 --- a/src/infrastructure/db.rs +++ b/src/infrastructure/db.rs @@ -846,6 +846,161 @@ fn run_migrations(conn: &Connection) -> Result<(), String> { log::info!("Database migration to version 17 completed"); } + // Migration v18: Add OAuth columns to user_settings + if current_version < 18 { + log::info!("Migrating database from version {} to 18", current_version); + + // Check and add auth_method column + let auth_method_exists = conn + .prepare( + "SELECT name FROM pragma_table_info('user_settings') WHERE name = 'auth_method'", + ) + .and_then(|mut stmt| { + stmt.query_map([], |row| row.get::<_, String>(0)) + .map(|mut rows| rows.next().is_some()) + }) + .unwrap_or(false); + + if !auth_method_exists { + log::info!("Adding auth_method column to user_settings table"); + conn.execute( + "ALTER TABLE user_settings ADD COLUMN auth_method TEXT NOT NULL DEFAULT 'api_key'", + [], + ) + .map_err(|e| format!("Failed to add auth_method column: {}", e))?; + } + + // Check and add oauth_access_token column + let oauth_access_token_exists = conn + .prepare("SELECT name FROM pragma_table_info('user_settings') WHERE name = 'oauth_access_token'") + .and_then(|mut stmt| { + stmt.query_map([], |row| row.get::<_, String>(0)) + .map(|mut rows| rows.next().is_some()) + }) + .unwrap_or(false); + + if !oauth_access_token_exists { + log::info!("Adding oauth_access_token column to user_settings table"); + conn.execute( + "ALTER TABLE user_settings ADD COLUMN oauth_access_token TEXT", + [], + ) + .map_err(|e| format!("Failed to add oauth_access_token column: {}", e))?; + } + + // Check and add oauth_refresh_token column + let oauth_refresh_token_exists = conn + .prepare("SELECT name FROM pragma_table_info('user_settings') WHERE name = 'oauth_refresh_token'") + .and_then(|mut stmt| { + stmt.query_map([], |row| row.get::<_, String>(0)) + .map(|mut rows| rows.next().is_some()) + }) + .unwrap_or(false); + + if !oauth_refresh_token_exists { + log::info!("Adding oauth_refresh_token column to user_settings table"); + conn.execute( + "ALTER TABLE user_settings ADD COLUMN oauth_refresh_token TEXT", + [], + ) + .map_err(|e| format!("Failed to add oauth_refresh_token column: {}", e))?; + } + + // Check and add oauth_token_expiry column + let oauth_token_expiry_exists = conn + .prepare("SELECT name FROM pragma_table_info('user_settings') WHERE name = 'oauth_token_expiry'") + .and_then(|mut stmt| { + stmt.query_map([], |row| row.get::<_, String>(0)) + .map(|mut rows| rows.next().is_some()) + }) + .unwrap_or(false); + + if !oauth_token_expiry_exists { + log::info!("Adding oauth_token_expiry column to user_settings table"); + conn.execute( + "ALTER TABLE user_settings ADD COLUMN oauth_token_expiry INTEGER", + [], + ) + .map_err(|e| format!("Failed to add oauth_token_expiry column: {}", e))?; + } + + // Check and add oauth_account_id column + let oauth_account_id_exists = conn + .prepare("SELECT name FROM pragma_table_info('user_settings') WHERE name = 'oauth_account_id'") + .and_then(|mut stmt| { + stmt.query_map([], |row| row.get::<_, String>(0)) + .map(|mut rows| rows.next().is_some()) + }) + .unwrap_or(false); + + if !oauth_account_id_exists { + log::info!("Adding oauth_account_id column to user_settings table"); + conn.execute( + "ALTER TABLE user_settings ADD COLUMN oauth_account_id TEXT", + [], + ) + .map_err(|e| format!("Failed to add oauth_account_id column: {}", e))?; + } + + conn.execute( + "INSERT OR REPLACE INTO meta (key, value) VALUES ('schema_version', '18')", + [], + ) + .map_err(|e| format!("Failed to update schema_version: {}", e))?; + + log::info!("Database migration to version 18 completed"); + } + + // Migration v19: add auth_type to models + if current_version < 19 { + log::info!("Migrating database from version {} to 19", current_version); + + let auth_type_exists = conn + .prepare("SELECT name FROM pragma_table_info('models') WHERE name = 'auth_type'") + .and_then(|mut stmt| { + stmt.query_map([], |row| row.get::<_, String>(0)) + .map(|mut rows| rows.next().is_some()) + }) + .unwrap_or(false); + + if !auth_type_exists { + log::info!("Adding auth_type column to models table"); + conn.execute( + "ALTER TABLE models ADD COLUMN auth_type TEXT NOT NULL DEFAULT 'local'", + [], + ) + .map_err(|e| format!("Failed to add auth_type column: {}", e))?; + } + + let openai_auth_type: String = conn + .query_row( + "SELECT auth_method FROM user_settings WHERE id = 1", + [], + |row| row.get(0), + ) + .unwrap_or_else(|_| "api_key".to_string()); + + conn.execute( + "UPDATE models SET auth_type = 'local' WHERE type = 'local'", + [], + ) + .map_err(|e| format!("Failed to backfill local auth_type: {}", e))?; + + conn.execute( + "UPDATE models SET auth_type = ? WHERE type = 'openai' AND (auth_type IS NULL OR auth_type = '')", + params![openai_auth_type], + ) + .map_err(|e| format!("Failed to backfill OpenAI auth_type: {}", e))?; + + conn.execute( + "INSERT OR REPLACE INTO meta (key, value) VALUES ('schema_version', '19')", + [], + ) + .map_err(|e| format!("Failed to update schema_version: {}", e))?; + + log::info!("Database migration to version 19 completed"); + } + conn.execute( "INSERT OR IGNORE INTO user_settings (id, openai_tracing_enabled, use_behavior_trees, web_search_enabled, max_tool_calls_per_request) VALUES (1, 0, 0, 0, 50)", [], diff --git a/src/infrastructure/event_bus.rs b/src/infrastructure/event_bus.rs index 3e80898..5b8ef5d 100644 --- a/src/infrastructure/event_bus.rs +++ b/src/infrastructure/event_bus.rs @@ -1,8 +1,8 @@ -use crate::domain::permissions::PermissionDecision; +use crate::domain::permissions::UserPermissionDecision; +use crate::domain::todo::TodoItem; use crate::domain::tools::FileChange; use crate::domain::AgentModeType; use crossbeam_channel::{unbounded, Receiver, Sender}; -use crate::domain::todo::TodoItem; #[derive(Debug, Clone)] pub enum StepPhase { @@ -33,6 +33,7 @@ pub struct LocalModelInfo { pub struct OpenAiModelInfo { pub _id: i64, pub name: String, + pub auth_type: String, } #[derive(Debug, Clone)] @@ -48,12 +49,13 @@ pub enum UiToAgentEvent { mode: AgentModeType, // NEW: Pass mode to agent session_id: Option, }, + CancelRequest, SessionContinueEvent { session_id: i64, }, PermissionUpdateEvent { request_id: u64, - decision: PermissionDecision, + decision: UserPermissionDecision, }, ShutdownEvent, SettingsUpdateEvent { @@ -109,7 +111,7 @@ pub enum AgentToUiEvent { #[derive(Debug, Clone)] pub struct PermissionUpdate { pub request_id: u64, - pub decision: PermissionDecision, + pub decision: UserPermissionDecision, } #[derive(Clone, Debug)] diff --git a/src/infrastructure/event_controller.rs b/src/infrastructure/event_controller.rs index 60b7e52..5421d22 100644 --- a/src/infrastructure/event_controller.rs +++ b/src/infrastructure/event_controller.rs @@ -1,13 +1,12 @@ -use crate::domain::permissions::{PermissionConfig, PermissionDecision, PermissionPrompter}; +use crate::domain::permissions::{PermissionConfig, PermissionPrompter, UserPermissionDecision}; use crate::domain::{CancellationToken, SessionService, StartupService}; use crate::infrastructure::db::DbPool; use crate::infrastructure::event_bus::{ - AgentToUiEvent, EventBus, ImageAttachment, PermissionUpdate, RequestStatus, StepPhase, UiToAgentEvent, + AgentToUiEvent, EventBus, ImageAttachment, PermissionUpdate, RequestStatus, StepPhase, + UiToAgentEvent, }; use crate::infrastructure::inference::InferenceEngine; -use crate::infrastructure::init::{ - apply_model_selection, update_openai_api_key, -}; +use crate::infrastructure::init::{apply_model_selection, update_openai_api_key}; use crate::repository::UserSettingsRepository; use crate::{domain, infrastructure}; use crossbeam_channel::{select, unbounded, Receiver, Sender}; @@ -18,11 +17,7 @@ use std::sync::Arc; use std::thread; /// Helper function to send a failure event to the UI -fn send_failure_event( - tx: &Sender, - request_id: u64, - error: impl Into, -) { +fn send_failure_event(tx: &Sender, request_id: u64, error: impl Into) { let _ = tx.send(AgentToUiEvent::RequestFinishedEvent { request_id, status: RequestStatus::Failure, @@ -85,7 +80,12 @@ pub struct EventController { cancel_token: CancellationToken, permission_response_tx: Option>, request_counter: u64, - pending_request: Option<(String, Vec, crate::domain::AgentModeType, Option)>, + pending_request: Option<( + String, + Vec, + crate::domain::AgentModeType, + Option, + )>, } impl EventController { @@ -188,20 +188,74 @@ impl EventController { } }; - if let Some(last_request) = requests.last() { - let label = request_label(&last_request.prompt); - let status = match &last_request.result_summary { + let steps_repo = crate::repository::SessionRequestStepsRepository::new(self.conn.clone()); + for request in &requests { + let label = request_label(&request.prompt); + let status = match &request.result_summary { Some(summary) if summary.starts_with("Success") => RequestStatus::Success, Some(_) => RequestStatus::Failure, None => RequestStatus::Failure, }; - let summary = last_request.result_summary.clone(); + let summary = request.result_summary.clone(); let _ = self.bus.agent_to_ui_tx.send(AgentToUiEvent::RequestStartedEvent { request_id: 0, label, - prompt: last_request.prompt.clone(), + prompt: request.prompt.clone(), }); + + let steps = match steps_repo.load_steps_for_request(request.id) { + Ok(steps) => steps, + Err(e) => { + log::warn!( + "Failed to load steps for request {}: {}", + request.id, + e + ); + Vec::new() + } + }; + + if steps.is_empty() { + if let Some(log) = &request._steps_log { + for line in log.lines() { + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + let _ = self.bus.agent_to_ui_tx.send(AgentToUiEvent::ProgressEvent { + step_name: "assistant_response".to_string(), + phase: StepPhase::Before, + summary: trimmed.to_string(), + }); + let _ = self.bus.agent_to_ui_tx.send(AgentToUiEvent::ProgressEvent { + step_name: "assistant_response".to_string(), + phase: StepPhase::After, + summary: trimmed.to_string(), + }); + } + } + } else { + for step in steps { + let step_name = step + .tool_name + .clone() + .unwrap_or_else(|| step.step_type.clone()); + let summary = step.summary.clone(); + + let _ = self.bus.agent_to_ui_tx.send(AgentToUiEvent::ProgressEvent { + step_name: step_name.clone(), + phase: StepPhase::Before, + summary: summary.clone(), + }); + let _ = self.bus.agent_to_ui_tx.send(AgentToUiEvent::ProgressEvent { + step_name, + phase: StepPhase::After, + summary, + }); + } + } + let _ = self.bus.agent_to_ui_tx.send(AgentToUiEvent::RequestFinishedEvent { request_id: 0, status, @@ -299,6 +353,12 @@ impl EventController { let _ = worker_status_tx.send(()); }); } + UiToAgentEvent::CancelRequest => { + if worker_running { + log::info!("Cancelling current request via Ctrl+C"); + self.cancel_token.cancel(); + } + } UiToAgentEvent::PermissionUpdateEvent { request_id, decision, @@ -332,7 +392,8 @@ impl EventController { match apply_model_selection(&self.conn, selection) { Ok(engine) => { self.engine = Some(engine); - self.current_session_id = None; + // Keep current session when switching models so ongoing + // conversation continues with a different model. } Err(err) => { send_failure_and_continue!(self, err.to_string()); @@ -443,7 +504,7 @@ fn run_request_worker( use_behavior_trees, PermissionConfig::default(), prompter, - event_bus.agent_to_ui_tx.clone() + event_bus.agent_to_ui_tx.clone(), ) .map_err(|e| format!("Failed to create session service: {}", e))?; @@ -462,32 +523,58 @@ fn run_request_worker( } else { chain.fail_reason }; - let _ = event_bus.agent_to_ui_tx.send(AgentToUiEvent::RequestFinishedEvent { - request_id, - status: RequestStatus::Failure, - summary: Some(reason), - final_message: None, - }); + let _ = event_bus + .agent_to_ui_tx + .send(AgentToUiEvent::RequestFinishedEvent { + request_id, + status: RequestStatus::Failure, + summary: Some(reason), + final_message: None, + }); } else { - let _ = event_bus.agent_to_ui_tx.send(AgentToUiEvent::RequestFinishedEvent { - request_id, - status: RequestStatus::Success, - summary: Some(chain.get_summary()), - final_message: chain.final_message.map(|msg| msg.to_string()), - }); + let _ = event_bus + .agent_to_ui_tx + .send(AgentToUiEvent::RequestFinishedEvent { + request_id, + status: RequestStatus::Success, + summary: Some(chain.get_summary()), + final_message: chain.final_message.map(|msg| msg.to_string()), + }); } Ok(()) } Err(domain::session::ServiceError::Workflow(domain::workflow::Error::Cancelled)) => { - let _ = event_bus.agent_to_ui_tx.send(AgentToUiEvent::RequestFinishedEvent { - request_id, - status: RequestStatus::Cancelled, - summary: Some("Cancelled".to_string()), - final_message: None, - }); + let _ = event_bus + .agent_to_ui_tx + .send(AgentToUiEvent::RequestFinishedEvent { + request_id, + status: RequestStatus::Cancelled, + summary: Some("Cancelled".to_string()), + final_message: None, + }); Ok(()) } - Err(err) => Err(format!("Workflow error: {}", err)), + Err(err) => { + let err_str = err.to_string(); + + // Check if error is due to expired OAuth token + if err_str.contains("AuthExpired") || err_str.contains("OAuth token expired") { + let _ = event_bus + .agent_to_ui_tx + .send(AgentToUiEvent::RequestFinishedEvent { + request_id, + status: RequestStatus::Failure, + summary: Some( + "OAuth token expired. Please run /oauth-login to re-authenticate." + .to_string(), + ), + final_message: None, + }); + return Ok(()); + } + + Err(format!("Workflow error: {}", err)) + } } } @@ -533,7 +620,7 @@ fn emit_chain_progress( #[cfg(test)] mod tests { use super::*; - use crate::domain::permissions::{PermissionDecision, PermissionRequest, PermissionScope}; + use crate::domain::permissions::{PermissionRequest, PermissionScope, UserPermissionDecision}; use crate::infrastructure::db; use crate::infrastructure::event_bus::{ModelSelection, UiToAgentEvent}; use crate::repository::{MetaRepository, ModelsRepository, UserSettingsRepository}; @@ -646,6 +733,7 @@ mod tests { assert!(model.is_some()); let model = model.unwrap(); assert!(model._api_key.is_none()); + assert_eq!(model.auth_type, domain::ModelAuthType::ApiKey); let settings_repo = UserSettingsRepository::new(&*conn_guard); let settings = settings_repo.get_current().expect("settings"); @@ -663,7 +751,7 @@ mod tests { vec![], PermissionScope::Project, Some(1), - false, // is_read_only + false, // is_read_only PathBuf::from("/tmp"), ); @@ -681,11 +769,11 @@ mod tests { let _ = update_tx.send(PermissionUpdate { request_id, - decision: PermissionDecision::AllowOnce, + decision: UserPermissionDecision::AllowOnce, }); let decision = handler.join().expect("join handler"); - assert_eq!(decision, PermissionDecision::AllowOnce); + assert_eq!(decision, UserPermissionDecision::AllowOnce); } } @@ -709,7 +797,7 @@ impl PermissionPrompter for BusPermissionPrompter { fn ask_permission( &self, request: &domain::permissions::PermissionRequest, - ) -> Result { + ) -> Result { let request_id = self.counter.fetch_add(1, Ordering::SeqCst); let scope = match request.scope { domain::permissions::PermissionScope::Session => "session", diff --git a/src/infrastructure/inference/openai.rs b/src/infrastructure/inference/openai.rs index 6761e29..aff9137 100644 --- a/src/infrastructure/inference/openai.rs +++ b/src/infrastructure/inference/openai.rs @@ -1,5 +1,6 @@ use super::{InferenceEngine, LLMInferenceResult}; use crate::domain::ModelType; +use crate::infrastructure::api_clients::openai::client::AuthToken; use crate::infrastructure::api_clients::openai::OpenAIClient; use crate::infrastructure::InfaError; use openai_api_rust::models::ModelsApi; @@ -12,11 +13,17 @@ pub struct OpenAIEngine { } impl OpenAIEngine { - /// Create a new OpenAI engine with the given API key and model - pub fn new(api_key: &str, model: &str) -> Result, String> { - let auth = Auth::new(api_key); + /// Create a new OpenAI engine with the given auth token and model + pub fn new(auth_token: AuthToken, model: &str) -> Result, String> { + // Extract API key string for openai_api_rust client (used for model listing) + let api_key_str = match &auth_token { + AuthToken::ApiKey(key) => key.clone(), + AuthToken::OAuth { token, .. } => token.clone(), + }; + + let auth = Auth::new(&api_key_str); let client = OpenAI::new(auth, "https://api.openai.com/v1/"); - let responses_client = OpenAIClient::new(api_key.to_string(), model.to_string()); + let responses_client = OpenAIClient::new(auth_token, model.to_string()); let engine = Arc::new(Self { client, @@ -26,11 +33,17 @@ impl OpenAIEngine { Ok(engine as Arc) } - /// Create a new OpenAI engine with the given API key and model - pub fn new_general(api_key: &str, model: &str) -> Result { - let auth = Auth::new(api_key); + /// Create a new OpenAI engine with the given auth token and model + pub fn new_general(auth_token: AuthToken, model: &str) -> Result { + // Extract API key string for openai_api_rust client (used for model listing) + let api_key_str = match &auth_token { + AuthToken::ApiKey(key) => key.clone(), + AuthToken::OAuth { token, .. } => token.clone(), + }; + + let auth = Auth::new(&api_key_str); let client = OpenAI::new(auth, "https://api.openai.com/v1/"); - let responses_client = OpenAIClient::new(api_key.to_string(), model.to_string()); + let responses_client = OpenAIClient::new(auth_token, model.to_string()); let engine = Self { client, diff --git a/src/infrastructure/init.rs b/src/infrastructure/init.rs index 4321763..ca609f0 100644 --- a/src/infrastructure/init.rs +++ b/src/infrastructure/init.rs @@ -1,4 +1,6 @@ -use crate::domain::ModelType; +use crate::domain::{ModelAuthType, ModelType}; +use crate::infrastructure::api_clients::openai::client::AuthToken; +use crate::infrastructure::auth::refresh_oauth_tokens; use crate::infrastructure::db; use crate::infrastructure::db::DbPool; use crate::infrastructure::event_bus::{LocalModelInfo, ModelSelection}; @@ -30,6 +32,12 @@ pub enum InitError { #[error("OpenAI model missing api_key")] MissingApiKey, + #[error("OAuth token missing")] + MissingOAuthToken, + + #[error("OAuth token expired - please re-authenticate")] + OAuthExpired, + #[error("failed to load model engine: {0}")] ModelLoadError(String), @@ -152,11 +160,63 @@ impl InfrastructureInitializer { ModelType::OpenAI => { let settings_repo = UserSettingsRepository::new(conn); let settings = settings_repo.get_current().map_err(InitError::Repository)?; - let api_key = settings.openai_api_key.ok_or(InitError::MissingApiKey)?; // Use saved model_name, fallback to gpt-4 for backward compatibility let model_name = model.model_name.as_deref().unwrap_or("gpt-4"); - use crate::infrastructure::inference::openai::OpenAIEngine; - OpenAIEngine::new(&api_key, model_name).map_err(|e| InitError::ModelLoadError(e)) + + let auth_token = match model.auth_type { + ModelAuthType::OAuth => { + let mut access_token = settings + .oauth_access_token + .ok_or(InitError::MissingOAuthToken)?; + let refresh_token_val = settings + .oauth_refresh_token + .ok_or(InitError::MissingOAuthToken)?; + let account_id = settings.oauth_account_id; + + // Check if token is expired and refresh if needed + if settings + .oauth_token_expiry + .map(|exp| chrono::Utc::now().timestamp() >= exp) + .unwrap_or(true) + { + log::info!("OAuth token expired at startup, refreshing..."); + let new_tokens = + refresh_oauth_tokens(&refresh_token_val).map_err(|e| { + InitError::ModelLoadError(format!( + "Token refresh failed: {}", + e + )) + })?; + + settings_repo + .update_oauth_tokens( + &new_tokens.access_token, + &new_tokens.refresh_token, + new_tokens.expires_in, + new_tokens.account_id.as_deref(), + ) + .map_err(InitError::Repository)?; + + access_token = new_tokens.access_token; + } + + AuthToken::OAuth { + token: access_token, + account_id, + } + } + ModelAuthType::ApiKey => { + let api_key = settings.openai_api_key.ok_or(InitError::MissingApiKey)?; + AuthToken::ApiKey(api_key) + } + ModelAuthType::Local => { + return Err(InitError::ModelLoadError( + "OpenAI model has local auth_type".to_string(), + )); + } + }; + + OpenAIEngine::new(auth_token, model_name).map_err(|e| InitError::ModelLoadError(e)) } } } @@ -290,6 +350,11 @@ pub fn apply_model_selection( .unwrap_or_else(|_| PathBuf::from(existing_path)); if existing_abs == canonical_path { model_id = Some(model.id); + if model.auth_type != ModelAuthType::Local { + models_repo + .update_auth_type(model.id, ModelAuthType::Local) + .map_err(InitError::Repository)?; + } break; } } @@ -299,7 +364,13 @@ pub fn apply_model_selection( Some(id) => id, None => { let model = models_repo - .create(ModelType::Local, None, Some(&gguf_path), None) + .create( + ModelType::Local, + None, + Some(&gguf_path), + None, + ModelAuthType::Local, + ) .map_err(InitError::Repository)?; model.id } @@ -324,12 +395,70 @@ pub fn apply_model_selection( .map_err(InitError::Repository)?; let settings = settings_repo.get_current().map_err(InitError::Repository)?; - let resolved_api_key = settings.openai_api_key.ok_or(InitError::MissingApiKey)?; + let auth_method = crate::domain::AuthMethod::from_str(&settings.auth_method); + let auth_type = ModelAuthType::from_auth_method(&auth_method); + + // Determine auth token based on auth method + let auth_token = match auth_type { + ModelAuthType::OAuth => { + let mut access_token = settings + .oauth_access_token + .ok_or(InitError::MissingOAuthToken)?; + let refresh_token_val = settings + .oauth_refresh_token + .ok_or(InitError::MissingOAuthToken)?; + let account_id = settings.oauth_account_id; + + // Check if token is expired and refresh if needed + if settings + .oauth_token_expiry + .map(|exp| chrono::Utc::now().timestamp() >= exp) + .unwrap_or(true) + { + log::info!("OAuth token expired, refreshing..."); + + let new_tokens = refresh_oauth_tokens(&refresh_token_val).map_err(|e| { + InitError::ModelLoadError(format!("Token refresh failed: {}", e)) + })?; + + // Update database with new tokens + settings_repo + .update_oauth_tokens( + &new_tokens.access_token, + &new_tokens.refresh_token, + new_tokens.expires_in, + new_tokens.account_id.as_deref(), + ) + .map_err(InitError::Repository)?; + + access_token = new_tokens.access_token; + } + + AuthToken::OAuth { + token: access_token, + account_id: account_id, + } + } + ModelAuthType::ApiKey => { + let key = settings.openai_api_key.ok_or(InitError::MissingApiKey)?; + AuthToken::ApiKey(key) + } + ModelAuthType::Local => { + return Err(InitError::ModelLoadError( + "OpenAI model has local auth_type".to_string(), + )); + } + }; let mut model_id = None; for model in &existing_models { if model.model_name.as_deref() == Some(&model_name) { model_id = Some(model.id); + if model.auth_type != auth_type { + models_repo + .update_auth_type(model.id, auth_type) + .map_err(InitError::Repository)?; + } break; } } @@ -342,6 +471,11 @@ pub fn apply_model_selection( models_repo .update_model_name(existing.id, Some(&model_name)) .map_err(InitError::Repository)?; + if existing.auth_type != auth_type { + models_repo + .update_auth_type(existing.id, auth_type) + .map_err(InitError::Repository)?; + } model_id = Some(existing.id); } } @@ -350,14 +484,14 @@ pub fn apply_model_selection( Some(id) => id, None => { models_repo - .create(ModelType::OpenAI, None, None, Some(&model_name)) + .create(ModelType::OpenAI, None, None, Some(&model_name), auth_type) .map_err(InitError::Repository)? .id } }; - let engine = OpenAIEngine::new(&resolved_api_key, &model_name) - .map_err(InitError::ModelLoadError)?; + let engine = + OpenAIEngine::new(auth_token, &model_name).map_err(InitError::ModelLoadError)?; meta_repo .set_last_used_model_id(model_id) diff --git a/src/infrastructure/mod.rs b/src/infrastructure/mod.rs index 50adff5..0f77b71 100644 --- a/src/infrastructure/mod.rs +++ b/src/infrastructure/mod.rs @@ -1,4 +1,5 @@ mod api_clients; +pub mod auth; pub mod cli; pub mod db; pub mod event_bus; diff --git a/src/repository/models.rs b/src/repository/models.rs index d67d5cd..ce1e72a 100644 --- a/src/repository/models.rs +++ b/src/repository/models.rs @@ -1,4 +1,4 @@ -use crate::domain::ModelType; +use crate::domain::{ModelAuthType, ModelType}; use rusqlite::{params, Connection, Row}; /// Raw database row for models table @@ -9,6 +9,7 @@ pub struct ModelRow { pub _api_key: Option, pub gguf_file_path: Option, pub model_name: Option, + pub auth_type: ModelAuthType, pub _date_added: String, } @@ -16,6 +17,17 @@ impl ModelRow { fn from_row(row: &Row) -> Result { let model_type_str: String = row.get(1)?; let model_type = ModelType::from_str(&model_type_str).unwrap_or(ModelType::Local); + let auth_type_str: Option = row.get(5)?; + let auth_type = match auth_type_str.as_deref() { + Some(value) => ModelAuthType::from_str(value), + None => { + if model_type == ModelType::Local { + ModelAuthType::Local + } else { + ModelAuthType::ApiKey + } + } + }; Ok(Self { id: row.get(0)?, @@ -23,7 +35,8 @@ impl ModelRow { _api_key: row.get(2)?, gguf_file_path: row.get(3)?, model_name: row.get(4)?, - _date_added: row.get(5)?, + auth_type, + _date_added: row.get(6)?, }) } } @@ -40,7 +53,7 @@ impl<'a> ModelsRepository<'a> { pub fn find_by_id(&self, id: i64) -> Result, String> { let mut stmt = self .conn - .prepare("SELECT id, type, api_key, gguf_file_path, model_name, date_added FROM models WHERE id = ?") + .prepare("SELECT id, type, api_key, gguf_file_path, model_name, auth_type, date_added FROM models WHERE id = ?") .map_err(|e| e.to_string())?; let result = stmt @@ -54,7 +67,7 @@ impl<'a> ModelsRepository<'a> { pub fn find_by_type(&self, model_type: ModelType) -> Result, String> { let mut stmt = self .conn - .prepare("SELECT id, type, api_key, gguf_file_path, model_name, date_added FROM models WHERE type = ? ORDER BY date_added DESC") + .prepare("SELECT id, type, api_key, gguf_file_path, model_name, auth_type, date_added FROM models WHERE type = ? ORDER BY date_added DESC") .map_err(|e| e.to_string())?; let rows = stmt @@ -75,13 +88,21 @@ impl<'a> ModelsRepository<'a> { api_key: Option<&str>, gguf_file_path: Option<&str>, model_name: Option<&str>, + auth_type: ModelAuthType, ) -> Result { let date_added = chrono_now(); self.conn .execute( - "INSERT INTO models (type, api_key, gguf_file_path, model_name, date_added) VALUES (?, ?, ?, ?, ?)", - params![model_type.as_str(), api_key, gguf_file_path, model_name, date_added], + "INSERT INTO models (type, api_key, gguf_file_path, model_name, auth_type, date_added) VALUES (?, ?, ?, ?, ?, ?)", + params![ + model_type.as_str(), + api_key, + gguf_file_path, + model_name, + auth_type.as_str(), + date_added + ], ) .map_err(|e| e.to_string())?; @@ -93,10 +114,23 @@ impl<'a> ModelsRepository<'a> { _api_key: api_key.map(|s| s.to_string()), gguf_file_path: gguf_file_path.map(|s| s.to_string()), model_name: model_name.map(|s| s.to_string()), + auth_type, _date_added: date_added, }) } + pub fn update_auth_type(&self, id: i64, auth_type: ModelAuthType) -> Result { + let rows_affected = self + .conn + .execute( + "UPDATE models SET auth_type = ? WHERE id = ?", + params![auth_type.as_str(), id], + ) + .map_err(|e| e.to_string())?; + + Ok(rows_affected > 0) + } + pub fn update_model_name(&self, id: i64, model_name: Option<&str>) -> Result { let rows_affected = self .conn diff --git a/src/repository/user_settings.rs b/src/repository/user_settings.rs index 21c9996..d371778 100644 --- a/src/repository/user_settings.rs +++ b/src/repository/user_settings.rs @@ -10,6 +10,11 @@ pub struct UserSettingsRow { pub web_search_enabled: bool, pub brave_api_key: Option, pub max_tool_calls_per_request: i32, + pub auth_method: String, + pub oauth_access_token: Option, + pub oauth_refresh_token: Option, + pub oauth_token_expiry: Option, + pub oauth_account_id: Option, } impl UserSettingsRow { @@ -23,6 +28,13 @@ impl UserSettingsRow { web_search_enabled: row.get::<_, i64>(5)? != 0, brave_api_key: row.get(6)?, max_tool_calls_per_request: row.get(7)?, + auth_method: row + .get::<_, Option>(8)? + .unwrap_or_else(|| "api_key".to_string()), + oauth_access_token: row.get(9)?, + oauth_refresh_token: row.get(10)?, + oauth_token_expiry: row.get(11)?, + oauth_account_id: row.get(12)?, }) } } @@ -46,7 +58,7 @@ impl<'a> UserSettingsRepository<'a> { let mut stmt = self .conn .prepare( - "SELECT id, openai_api_key, openai_tracing_enabled, use_behavior_trees, current_model_id, web_search_enabled, brave_api_key, max_tool_calls_per_request FROM user_settings WHERE id = ?", + "SELECT id, openai_api_key, openai_tracing_enabled, use_behavior_trees, current_model_id, web_search_enabled, brave_api_key, max_tool_calls_per_request, auth_method, oauth_access_token, oauth_refresh_token, oauth_token_expiry, oauth_account_id FROM user_settings WHERE id = ?", ) .map_err(|e| e.to_string())?; @@ -62,7 +74,7 @@ impl<'a> UserSettingsRepository<'a> { self.ensure_row()?; self.conn .execute( - "UPDATE user_settings SET openai_api_key = ? WHERE id = 1", + "UPDATE user_settings SET openai_api_key = ?, auth_method = 'api_key' WHERE id = 1", params![api_key], ) .map_err(|e| e.to_string())?; @@ -114,6 +126,49 @@ impl<'a> UserSettingsRepository<'a> { Ok(()) } + pub fn update_oauth_tokens( + &self, + access_token: &str, + refresh_token: &str, + expires_in: i64, + account_id: Option<&str>, + ) -> Result<(), String> { + self.ensure_row()?; + let expiry = chrono::Utc::now().timestamp() + expires_in; + self.conn + .execute( + "UPDATE user_settings SET + oauth_access_token = ?, + oauth_refresh_token = ?, + oauth_token_expiry = ?, + oauth_account_id = ?, + auth_method = 'oauth', + openai_api_key = NULL + WHERE id = 1", + params![access_token, refresh_token, expiry, account_id], + ) + .map_err(|e| e.to_string())?; + Ok(()) + } + + #[allow(dead_code)] + pub fn clear_oauth_tokens(&self) -> Result<(), String> { + self.ensure_row()?; + self.conn + .execute( + "UPDATE user_settings SET + oauth_access_token = NULL, + oauth_refresh_token = NULL, + oauth_token_expiry = NULL, + oauth_account_id = NULL, + auth_method = 'api_key' + WHERE id = 1", + [], + ) + .map_err(|e| e.to_string())?; + Ok(()) + } + fn update_bool(&self, column: &str, enabled: bool) -> Result<(), String> { self.ensure_row()?; let sql = format!("UPDATE user_settings SET {} = ? WHERE id = 1", column);