├── .github └── workflows │ ├── build.yml │ └── release.yml ├── .gitignore ├── Cargo.lock ├── Cargo.toml ├── LICENSE ├── README.md └── crates ├── code_assistant ├── Cargo.lock ├── Cargo.toml ├── assets │ └── icons │ │ ├── LICENSES │ │ ├── ai_anthropic.svg │ │ ├── ai_google.svg │ │ ├── ai_ollama.svg │ │ ├── ai_open_ai.svg │ │ ├── arrow_circle.svg │ │ ├── arrow_down.svg │ │ ├── arrow_left.svg │ │ ├── arrow_right.svg │ │ ├── arrow_up.svg │ │ ├── brain.svg │ │ ├── caret_down.svg │ │ ├── caret_up.svg │ │ ├── check.svg │ │ ├── check_circle.svg │ │ ├── chevron_down.svg │ │ ├── chevron_down_small.svg │ │ ├── chevron_left.svg │ │ ├── chevron_right.svg │ │ ├── chevron_up.svg │ │ ├── chevron_up_down.svg │ │ ├── circle_stop.svg │ │ ├── close.svg │ │ ├── code.svg │ │ ├── delete.svg │ │ ├── ellipsis.svg │ │ ├── ellipsis_vertical.svg │ │ ├── exit.svg │ │ ├── expand_vertical.svg │ │ ├── file_code.svg │ │ ├── file_generic.svg │ │ ├── file_icons │ │ ├── archive.svg │ │ ├── astro.svg │ │ ├── audio.svg │ │ ├── book.svg │ │ ├── bun.svg │ │ ├── c.svg │ │ ├── camera.svg │ │ ├── code.svg │ │ ├── coffeescript.svg │ │ ├── conversations.svg │ │ ├── cpp.svg │ │ ├── css.svg │ │ ├── dart.svg │ │ ├── database.svg │ │ ├── docker.svg │ │ ├── elixir.svg │ │ ├── elm.svg │ │ ├── erlang.svg │ │ ├── eslint.svg │ │ ├── file.svg │ │ ├── file_types.json │ │ ├── folder.svg │ │ ├── folder_open.svg │ │ ├── font.svg │ │ ├── fsharp.svg │ │ ├── git.svg │ │ ├── go.svg │ │ ├── graphql.svg │ │ ├── hash.svg │ │ ├── haskell.svg │ │ ├── heroku.svg │ │ ├── html.svg │ │ ├── image.svg │ │ ├── info.svg │ │ ├── java.svg │ │ ├── javascript.svg │ │ ├── kotlin.svg │ │ ├── lock.svg │ │ ├── lua.svg │ │ ├── magnifying_glass.svg │ │ ├── nim.svg │ │ ├── notebook.svg │ │ ├── ocaml.svg │ │ ├── package.svg │ │ ├── phoenix.svg │ │ ├── php.svg │ │ ├── prettier.svg │ │ ├── prisma.svg │ │ ├── project.svg │ │ ├── python.svg │ │ ├── r.svg │ │ ├── react.svg │ │ ├── ruby.svg │ │ ├── rust.svg │ │ ├── scala.svg │ │ ├── swift.svg │ │ ├── tcl.svg │ │ ├── terraform.svg │ │ ├── toml.svg │ │ ├── typescript.svg │ │ ├── video.svg │ │ └── vue.svg │ │ ├── file_tree.svg │ │ ├── generic_close.svg │ │ ├── generic_maximize.svg │ │ ├── generic_minimize.svg │ │ ├── generic_restore.svg │ │ ├── history_rerun.svg │ │ ├── library.svg │ │ ├── link.svg │ │ ├── list_tree.svg │ │ ├── magnifying_glass.svg │ │ ├── maximize.svg │ │ ├── menu.svg │ │ ├── message_bubbles.svg │ │ ├── minimize.svg │ │ ├── panel_right_close.svg │ │ ├── panel_right_open.svg │ │ ├── pencil.svg │ │ ├── person.svg │ │ ├── plus.svg │ │ ├── replace.svg │ │ ├── replace_all.svg │ │ ├── replace_next.svg │ │ ├── rerun.svg │ │ ├── return.svg │ │ ├── reveal.svg │ │ ├── rotate_ccw.svg │ │ ├── rotate_cw.svg │ │ ├── search_code.svg │ │ ├── send.svg │ │ ├── settings.svg │ │ ├── settings_alt.svg │ │ ├── stop.svg │ │ ├── terminal.svg │ │ ├── text_snippet.svg │ │ ├── theme_dark.svg │ │ ├── theme_light.svg │ │ └── trash.svg ├── resources │ ├── system_message.md │ └── system_message_tools.md └── src │ ├── agent │ ├── mod.rs │ ├── runner.rs │ ├── tests.rs │ ├── tool_description_generator.rs │ └── types.rs │ ├── config.rs │ ├── explorer.rs │ ├── main.rs │ ├── mcp │ ├── handler.rs │ ├── mod.rs │ ├── resources.rs │ ├── server.rs │ ├── tests.rs │ └── types.rs │ ├── persistence.rs │ ├── tests │ ├── gitignore_tests.rs │ ├── mocks.rs │ └── mod.rs │ ├── tools │ ├── core │ │ ├── dyn_tool.rs │ │ ├── mod.rs │ │ ├── registry.rs │ │ ├── render.rs │ │ ├── result.rs │ │ ├── spec.rs │ │ └── tool.rs │ ├── impls │ │ ├── delete_files.rs │ │ ├── execute_command.rs │ │ ├── list_files.rs │ │ ├── list_projects.rs │ │ ├── mod.rs │ │ ├── perplexity_ask.rs │ │ ├── read_files.rs │ │ ├── replace_in_file.rs │ │ ├── search_files.rs │ │ ├── web_fetch.rs │ │ ├── web_search.rs │ │ └── write_file.rs │ ├── mod.rs │ ├── parse.rs │ ├── tests.rs │ └── types.rs │ ├── types.rs │ ├── ui │ ├── gpui │ │ ├── assets.rs │ │ ├── content_renderer.rs │ │ ├── diff_renderer.rs │ │ ├── elements.rs │ │ ├── file_icons.rs │ │ ├── memory.rs │ │ ├── messages.rs │ │ ├── mod.rs │ │ ├── parameter_renderers.rs │ │ ├── path_util.rs │ │ ├── root.rs │ │ ├── simple_renderers.rs │ │ ├── theme.rs │ │ └── ui_events.rs │ ├── mod.rs │ ├── streaming │ │ ├── json_processor.rs │ │ ├── json_processor_tests.rs │ │ ├── mod.rs │ │ ├── test_utils.rs │ │ ├── xml_processor.rs │ │ └── xml_processor_tests.rs │ ├── terminal.rs │ └── terminal_test.rs │ └── utils │ ├── command.rs │ ├── encoding.rs │ ├── file_updater.rs │ ├── mod.rs │ └── writer.rs ├── llm ├── Cargo.toml └── src │ ├── aicore_converse.rs │ ├── aicore_invoke.rs │ ├── anthropic.rs │ ├── anthropic_playback.rs │ ├── auth.rs │ ├── config.rs │ ├── display.rs │ ├── lib.rs │ ├── ollama.rs │ ├── openai.rs │ ├── openrouter.rs │ ├── recording.rs │ ├── tests.rs │ ├── types.rs │ ├── utils.rs │ └── vertex.rs └── web ├── Cargo.toml └── src ├── client.rs ├── lib.rs ├── perplexity.rs └── tests.rs /.github/workflows/build.yml: -------------------------------------------------------------------------------- 1 | name: Build 2 | 3 | on: 4 | push: 5 | branches: ["main"] 6 | pull_request: 7 | branches: ["main"] 8 | 9 | env: 10 | CARGO_TERM_COLOR: always 11 | 12 | jobs: 13 | build: 14 | runs-on: macos-latest 15 | 16 | steps: 17 | - uses: actions/checkout@v4 18 | # - name: Install dependencies 19 | # run: | 20 | # sudo apt-get update 21 | # sudo apt-get install -y libxcb1-dev libxcb-shape0-dev libxcb-xfixes0-dev libxkbcommon-dev libxkbcommon-x11-dev libx11-dev 22 | - name: Check formatting 23 | run: cargo fmt --all -- --check 24 | # - name: Lint 25 | # run: cargo clippy --all-targets --all-features -- -D warnings 26 | - name: Build 27 | run: cargo build --verbose 28 | - name: Run tests 29 | run: cargo test --verbose 30 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | /sessions 3 | mcp-config.json 4 | .code-assistant.state.json 5 | tarpaulin-report.html 6 | .DS_Store 7 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [workspace] 2 | members = ["crates/code_assistant", "crates/llm"] 3 | 4 | resolver = "2" 5 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2025 Stephan Aßmus 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /crates/code_assistant/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "code-assistant" 3 | version = "0.1.6" 4 | edition = "2021" 5 | 6 | [dependencies] 7 | llm = { path = "../llm" } 8 | web = { path = "../web" } 9 | 10 | glob = "0.3" 11 | ignore = "0.4" 12 | walkdir = "2.5" 13 | percent-encoding = "2.3" 14 | tokio = { version = "1.44", features = ["full"] } 15 | tempfile = "3.18" 16 | 17 | # Terminal UI 18 | rustyline = "12.0.0" 19 | crossterm = "0.27.0" 20 | 21 | # GPUI related 22 | gpui = { git = "https://github.com/huacnlee/zed.git", branch = "webview" } 23 | gpui-component = { git = "https://github.com/longbridge/gpui-component.git" } 24 | smallvec = "1.14" 25 | rust-embed = { version = "8.4", features = ["include-exclude"] } 26 | 27 | # JSON (de)serialization 28 | serde = { version = "1.0", features = ["derive"] } 29 | serde_json = "1.0" 30 | 31 | # Error handling 32 | anyhow = "1.0" 33 | thiserror = "1.0" 34 | regex = "1.11" 35 | 36 | # Logging 37 | tracing = "0.1" 38 | tracing-subscriber = { version = "0.3", features = ["env-filter"] } 39 | 40 | # CLI 41 | clap = { version = "4.5", features = ["derive"] } 42 | 43 | async-trait = "0.1" 44 | dotenv = "0.15" 45 | dirs = "5.0" 46 | md5 = "0.7.0" 47 | 48 | # Date and time handling 49 | chrono = { version = "0.4", features = ["serde"] } 50 | 51 | # File content inspection 52 | content_inspector = "0.2" 53 | encoding_rs = "0.8.35" 54 | unicode-segmentation = "1.12.0" 55 | rand = "0.8.5" 56 | 57 | # Diff visualization 58 | similar = { version = "2.5.0", features = ["inline"] } 59 | async-channel = "2.3.1" 60 | 61 | [dev-dependencies] 62 | axum = "0.7" 63 | bytes = "1.10" 64 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/LICENSES: -------------------------------------------------------------------------------- 1 | Lucide License 2 | 3 | ISC License 4 | 5 | Copyright (c) for portions of Lucide are held by Cole Bemis 2013-2022 as part of Feather (MIT). All other copyright (c) for Lucide are held by Lucide Contributors 2022. 6 | 7 | Permission to use, copy, modify, and/or distribute this software for any purpose with or without fee is hereby granted, provided that the above copyright notice and this permission notice appear in all copies. 8 | 9 | THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. 10 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/ai_anthropic.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/ai_google.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/ai_open_ai.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/arrow_circle.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/arrow_down.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/arrow_left.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/arrow_right.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/arrow_up.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/brain.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/caret_down.svg: -------------------------------------------------------------------------------- 1 | 2 | 8 | 9 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/caret_up.svg: -------------------------------------------------------------------------------- 1 | 2 | 8 | 9 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/check.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/check_circle.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/chevron_down.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/chevron_down_small.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/chevron_left.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/chevron_right.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/chevron_up.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/chevron_up_down.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/circle_stop.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/close.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/code.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/delete.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/ellipsis.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/ellipsis_vertical.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/exit.svg: -------------------------------------------------------------------------------- 1 | 2 | 8 | 9 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/expand_vertical.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_code.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_generic.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/archive.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/astro.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/audio.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/book.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/bun.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/c.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/camera.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/code.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/coffeescript.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/conversations.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/cpp.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/css.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/dart.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/database.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/docker.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/elixir.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/elm.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/erlang.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/eslint.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/file.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/folder.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/folder_open.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/font.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/fsharp.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/git.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/go.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/graphql.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/hash.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/haskell.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/heroku.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/html.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/image.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/info.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/java.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/javascript.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/kotlin.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/lock.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/lua.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/magnifying_glass.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/nim.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/notebook.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/ocaml.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/package.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/phoenix.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/php.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/prettier.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/prisma.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/project.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/python.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/r.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/react.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/ruby.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/rust.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/scala.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/swift.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/tcl.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/terraform.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/toml.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/typescript.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/video.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_icons/vue.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/file_tree.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/generic_close.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/generic_maximize.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/generic_minimize.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/generic_restore.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/history_rerun.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/library.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/link.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/list_tree.svg: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/magnifying_glass.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/maximize.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/menu.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/message_bubbles.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/minimize.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/panel_right_close.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/panel_right_open.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/pencil.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/person.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/plus.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/replace.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/replace_all.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/replace_next.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/rerun.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/return.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/reveal.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/rotate_ccw.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/rotate_cw.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/search_code.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/send.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/settings.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/settings_alt.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/stop.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/terminal.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/text_snippet.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/theme_dark.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/theme_light.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /crates/code_assistant/assets/icons/trash.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /crates/code_assistant/src/agent/mod.rs: -------------------------------------------------------------------------------- 1 | #[cfg(test)] 2 | mod tests; 3 | 4 | mod runner; 5 | mod tool_description_generator; 6 | mod types; 7 | 8 | pub use crate::types::ToolMode; 9 | pub use runner::Agent; 10 | -------------------------------------------------------------------------------- /crates/code_assistant/src/agent/types.rs: -------------------------------------------------------------------------------- 1 | use crate::tools::core::AnyOutput; 2 | use serde_json::Value; 3 | 4 | /// Represents a tool request from the LLM, derived from ContentBlock::ToolUse 5 | #[derive(Debug, Clone)] 6 | pub struct ToolRequest { 7 | pub id: String, 8 | pub name: String, 9 | pub input: Value, 10 | } 11 | 12 | impl From<&llm::ContentBlock> for ToolRequest { 13 | fn from(block: &llm::ContentBlock) -> Self { 14 | if let llm::ContentBlock::ToolUse { id, name, input } = block { 15 | Self { 16 | id: id.clone(), 17 | name: name.clone(), 18 | input: input.clone(), 19 | } 20 | } else { 21 | panic!("Cannot convert non-ToolUse ContentBlock to ToolRequest") 22 | } 23 | } 24 | } 25 | 26 | /// Record of a tool execution with its result 27 | pub struct ToolExecution { 28 | pub tool_request: ToolRequest, 29 | pub result: Box, 30 | } 31 | -------------------------------------------------------------------------------- /crates/code_assistant/src/config.rs: -------------------------------------------------------------------------------- 1 | use crate::explorer::Explorer; 2 | use crate::types::{CodeExplorer, Project}; 3 | use anyhow::Result; 4 | use std::collections::HashMap; 5 | use std::path::PathBuf; 6 | 7 | /// Get the path to the configuration file 8 | pub fn get_config_path() -> Result { 9 | let home = 10 | dirs::home_dir().ok_or_else(|| anyhow::anyhow!("Could not determine home directory"))?; 11 | let config_dir = home.join(".config").join("code-assistant"); 12 | std::fs::create_dir_all(&config_dir)?; // Ensure directory exists 13 | Ok(config_dir.join("projects.json")) 14 | } 15 | 16 | // The main trait for project management 17 | pub trait ProjectManager: Send + Sync { 18 | // Add a temporary project, returns the project name 19 | fn add_temporary_project(&mut self, path: PathBuf) -> Result; 20 | // Get all projects (both configured and temporary) 21 | fn get_projects(&self) -> Result>; 22 | fn get_project(&self, name: &str) -> Result>; 23 | fn get_explorer_for_project(&self, name: &str) -> Result>; 24 | } 25 | 26 | // Default implementation of ProjectManager that loads from config file 27 | pub struct DefaultProjectManager { 28 | temp_projects: HashMap, 29 | } 30 | 31 | impl DefaultProjectManager { 32 | pub fn new() -> Self { 33 | Self { 34 | temp_projects: HashMap::new(), 35 | } 36 | } 37 | } 38 | 39 | impl ProjectManager for DefaultProjectManager { 40 | fn add_temporary_project(&mut self, path: PathBuf) -> Result { 41 | // Canonicalize path 42 | let path = path.canonicalize()?; 43 | 44 | // Check if this path matches any existing project 45 | let projects = load_projects()?; 46 | for (name, project) in &projects { 47 | if project.path == path { 48 | return Ok(name.clone()); 49 | } 50 | } 51 | 52 | // Generate name from path leaf 53 | let mut name = path 54 | .file_name() 55 | .and_then(|n| n.to_str()) 56 | .unwrap_or("temp_project") 57 | .to_string(); 58 | 59 | // Ensure name is unique 60 | let mut counter = 1; 61 | let original_name = name.clone(); 62 | while projects.contains_key(&name) { 63 | name = format!("{}_{}", original_name, counter); 64 | counter += 1; 65 | } 66 | 67 | // Add to temporary projects 68 | self.temp_projects.insert(name.clone(), Project { path }); 69 | 70 | Ok(name) 71 | } 72 | 73 | fn get_projects(&self) -> Result> { 74 | let mut all_projects = load_projects()?; 75 | all_projects.extend(self.temp_projects.clone()); 76 | Ok(all_projects) 77 | } 78 | 79 | fn get_project(&self, name: &str) -> Result> { 80 | let projects = self.get_projects()?; 81 | Ok(projects.get(name).cloned()) 82 | } 83 | 84 | fn get_explorer_for_project(&self, name: &str) -> Result> { 85 | let project = self 86 | .get_project(name)? 87 | .ok_or_else(|| anyhow::anyhow!("Project not found: {}", name))?; 88 | 89 | Ok(Box::new(Explorer::new(project.path))) 90 | } 91 | } 92 | 93 | /// Load projects configuration from disk 94 | pub fn load_projects() -> Result> { 95 | let config_path = get_config_path()?; 96 | 97 | if !config_path.exists() { 98 | return Ok(HashMap::new()); 99 | } 100 | 101 | let content = std::fs::read_to_string(config_path)?; 102 | Ok(serde_json::from_str(&content)?) 103 | } 104 | -------------------------------------------------------------------------------- /crates/code_assistant/src/mcp/mod.rs: -------------------------------------------------------------------------------- 1 | mod handler; 2 | mod resources; 3 | mod server; 4 | mod types; 5 | 6 | #[cfg(test)] 7 | mod tests; 8 | 9 | pub use server::MCPServer; 10 | -------------------------------------------------------------------------------- /crates/code_assistant/src/mcp/resources.rs: -------------------------------------------------------------------------------- 1 | use super::types::{Resource, ResourceContent}; 2 | use crate::types::FileTreeEntry; 3 | use std::collections::HashSet; 4 | 5 | pub struct ResourceManager { 6 | file_tree: Option, 7 | subscriptions: HashSet, 8 | } 9 | 10 | impl ResourceManager { 11 | pub fn new() -> Self { 12 | Self { 13 | file_tree: None, 14 | subscriptions: HashSet::new(), 15 | } 16 | } 17 | 18 | /// Lists all available resources 19 | pub fn list_resources(&self) -> Vec { 20 | let mut resources = Vec::new(); 21 | // Add file tree resource if available 22 | if self.file_tree.is_some() { 23 | resources.push(Resource { 24 | uri: "tree:///".to_string(), 25 | name: "Repository Structure".to_string(), 26 | description: Some("The repository file tree structure".to_string()), 27 | mime_type: Some("text/plain".to_string()), 28 | }); 29 | } 30 | resources 31 | } 32 | 33 | /// Reads a specific resource content 34 | pub fn read_resource(&self, uri: &str) -> Option { 35 | match uri { 36 | "tree:///" => self.file_tree.as_ref().map(|t| ResourceContent { 37 | uri: uri.to_string(), 38 | mime_type: Some("text/plain".to_string()), 39 | text: Some(t.to_string()), 40 | }), 41 | _ => None, 42 | } 43 | } 44 | 45 | /// Subscribes to a resource 46 | pub fn subscribe(&mut self, uri: &str) { 47 | self.subscriptions.insert(uri.to_string()); 48 | } 49 | 50 | /// Unsubscribes from a resource 51 | pub fn unsubscribe(&mut self, uri: &str) { 52 | self.subscriptions.remove(uri); 53 | } 54 | 55 | /// Checks if a resource is subscribed 56 | pub fn is_subscribed(&self, uri: &str) -> bool { 57 | self.subscriptions.contains(uri) 58 | } 59 | 60 | /// Updates the file tree 61 | #[allow(dead_code)] 62 | pub fn update_file_tree(&mut self, tree: FileTreeEntry) { 63 | self.file_tree = Some(tree); 64 | } 65 | } 66 | -------------------------------------------------------------------------------- /crates/code_assistant/src/mcp/server.rs: -------------------------------------------------------------------------------- 1 | use crate::mcp::handler::MessageHandler; 2 | use anyhow::Result; 3 | use tokio::io::{stdin, AsyncBufReadExt, BufReader}; 4 | use tracing::{debug, error, trace}; 5 | 6 | pub struct MCPServer { 7 | handler: MessageHandler, 8 | } 9 | 10 | impl MCPServer { 11 | pub fn new() -> Result { 12 | Ok(Self { 13 | handler: MessageHandler::new(tokio::io::stdout())?, 14 | }) 15 | } 16 | 17 | pub async fn run(&mut self) -> Result<()> { 18 | debug!("Starting MCP server using stdio transport"); 19 | 20 | let stdin = stdin(); 21 | let mut reader = BufReader::new(stdin); 22 | 23 | let mut line = String::new(); 24 | while let Ok(n) = reader.read_line(&mut line).await { 25 | if n == 0 { 26 | break; // EOF 27 | } 28 | 29 | let trimmed = line.trim(); 30 | trace!("Received message: {}", trimmed); 31 | 32 | // Process the message 33 | match self.handler.handle_message(trimmed).await { 34 | Ok(()) => { 35 | trace!("Message processed successfully"); 36 | } 37 | Err(e) => { 38 | error!("Error handling message: {}", e); 39 | } 40 | } 41 | 42 | line.clear(); 43 | } 44 | 45 | debug!("MCP server shutting down"); 46 | Ok(()) 47 | } 48 | } 49 | -------------------------------------------------------------------------------- /crates/code_assistant/src/persistence.rs: -------------------------------------------------------------------------------- 1 | use anyhow::Result; 2 | use llm::Message; 3 | use serde::{Deserialize, Serialize}; 4 | use std::path::PathBuf; 5 | use tracing::debug; 6 | 7 | /// Persistent state of the agent 8 | #[derive(Debug, Serialize, Deserialize, Clone)] 9 | pub struct AgentState { 10 | /// Original task description 11 | pub task: String, 12 | /// Message history 13 | pub messages: Vec, 14 | } 15 | 16 | pub trait StatePersistence: Send + Sync { 17 | fn save_state(&mut self, task: String, messages: Vec) -> Result<()>; 18 | fn load_state(&mut self) -> Result>; 19 | fn cleanup(&mut self) -> Result<()>; 20 | } 21 | 22 | pub struct FileStatePersistence { 23 | root_dir: PathBuf, 24 | } 25 | 26 | impl FileStatePersistence { 27 | pub fn new(root_dir: PathBuf) -> Self { 28 | Self { root_dir } 29 | } 30 | } 31 | 32 | const STATE_FILE: &str = ".code-assistant.state.json"; 33 | 34 | impl StatePersistence for FileStatePersistence { 35 | fn save_state(&mut self, task: String, messages: Vec) -> Result<()> { 36 | let state = AgentState { task, messages }; 37 | let state_path = self.root_dir.join(STATE_FILE); 38 | debug!("Saving state to {}", state_path.display()); 39 | let json = serde_json::to_string_pretty(&state)?; 40 | std::fs::write(state_path, json)?; 41 | Ok(()) 42 | } 43 | 44 | fn load_state(&mut self) -> Result> { 45 | let state_path = self.root_dir.join(STATE_FILE); 46 | if !state_path.exists() { 47 | return Ok(None); 48 | } 49 | 50 | debug!("Loading state from {}", state_path.display()); 51 | let json = std::fs::read_to_string(state_path)?; 52 | let state = serde_json::from_str(&json)?; 53 | Ok(Some(state)) 54 | } 55 | 56 | fn cleanup(&mut self) -> Result<()> { 57 | let state_path = self.root_dir.join(STATE_FILE); 58 | if state_path.exists() { 59 | debug!("Removing state file {}", state_path.display()); 60 | std::fs::remove_file(state_path)?; 61 | } 62 | Ok(()) 63 | } 64 | } 65 | 66 | #[cfg(test)] 67 | pub struct MockStatePersistence { 68 | state: Option, 69 | } 70 | 71 | #[cfg(test)] 72 | impl MockStatePersistence { 73 | pub fn new() -> Self { 74 | Self { state: None } 75 | } 76 | } 77 | 78 | #[cfg(test)] 79 | impl StatePersistence for MockStatePersistence { 80 | fn save_state(&mut self, task: String, messages: Vec) -> Result<()> { 81 | // In-Memory state 82 | let state = AgentState { task, messages }; 83 | self.state = Some(state); 84 | Ok(()) 85 | } 86 | 87 | fn load_state(&mut self) -> Result> { 88 | Ok(self.state.clone()) 89 | } 90 | 91 | fn cleanup(&mut self) -> Result<()> { 92 | self.state = None; 93 | Ok(()) 94 | } 95 | } 96 | -------------------------------------------------------------------------------- /crates/code_assistant/src/tests/mod.rs: -------------------------------------------------------------------------------- 1 | pub mod gitignore_tests; 2 | pub mod mocks; 3 | -------------------------------------------------------------------------------- /crates/code_assistant/src/tools/core/dyn_tool.rs: -------------------------------------------------------------------------------- 1 | use super::render::Render; 2 | use super::result::ToolResult; 3 | use super::spec::ToolSpec; 4 | use super::tool::{Tool, ToolContext}; 5 | use crate::types::ToolError; 6 | use anyhow::Result; 7 | use serde::de::DeserializeOwned; 8 | use serde::{Deserialize, Serialize}; 9 | use serde_json::Value; 10 | 11 | /// Type-erased tool output that can be rendered and determined for success 12 | pub trait AnyOutput: Send + Sync { 13 | /// Get a reference to the output as a Render trait object 14 | fn as_render(&self) -> &dyn Render; 15 | 16 | /// Determine if the tool execution was successful 17 | fn is_success(&self) -> bool; 18 | 19 | /// Serialize this output to a JSON value 20 | #[allow(dead_code)] 21 | fn to_json(&self) -> Result; 22 | } 23 | 24 | /// Automatically implemented for all types that implement both Render, ToolResult and Serialize 25 | impl AnyOutput for T { 26 | fn as_render(&self) -> &dyn Render { 27 | self 28 | } 29 | 30 | fn is_success(&self) -> bool { 31 | ToolResult::is_success(self) 32 | } 33 | 34 | fn to_json(&self) -> Result { 35 | serde_json::to_value(self).map_err(|e| anyhow::anyhow!("Failed to serialize output: {}", e)) 36 | } 37 | } 38 | 39 | /// Type-erased tool interface for storing heterogeneous tools in collections 40 | #[async_trait::async_trait] 41 | pub trait DynTool: Send + Sync + 'static { 42 | /// Get the static metadata for this tool 43 | fn spec(&self) -> ToolSpec; 44 | 45 | /// Invoke the tool with JSON parameters and get a type-erased output 46 | async fn invoke<'a>( 47 | &self, 48 | context: &mut ToolContext<'a>, 49 | params: Value, 50 | ) -> Result>; 51 | 52 | /// Deserialize a JSON value into this tool's output type 53 | #[allow(dead_code)] 54 | fn deserialize_output(&self, json: Value) -> Result>; 55 | } 56 | 57 | /// Automatic implementation of DynTool for any type that implements Tool 58 | #[async_trait::async_trait] 59 | impl DynTool for T 60 | where 61 | T: Tool, 62 | T::Input: DeserializeOwned, 63 | T::Output: Render + ToolResult + Serialize + for<'de> Deserialize<'de> + Send + Sync + 'static, 64 | { 65 | fn spec(&self) -> ToolSpec { 66 | Tool::spec(self) 67 | } 68 | 69 | async fn invoke<'a>( 70 | &self, 71 | context: &mut ToolContext<'a>, 72 | params: Value, 73 | ) -> Result> { 74 | // Deserialize input 75 | let input: T::Input = serde_json::from_value(params).map_err(|e| { 76 | // Convert Serde error to ToolError::ParseError 77 | ToolError::ParseError(format!("Failed to parse parameters: {}", e)) 78 | })?; 79 | 80 | // Execute the tool 81 | let output = self.execute(context, input).await?; 82 | 83 | // Box the output as AnyOutput 84 | Ok(Box::new(output) as Box) 85 | } 86 | 87 | fn deserialize_output(&self, json: Value) -> Result> { 88 | // Use the tool's deserialize_output method 89 | let output = Tool::deserialize_output(self, json)?; 90 | 91 | // Box the output as AnyOutput 92 | Ok(Box::new(output) as Box) 93 | } 94 | } 95 | -------------------------------------------------------------------------------- /crates/code_assistant/src/tools/core/mod.rs: -------------------------------------------------------------------------------- 1 | // Core tools implementation 2 | pub mod dyn_tool; 3 | pub mod registry; 4 | pub mod render; 5 | pub mod result; 6 | pub mod spec; 7 | pub mod tool; 8 | 9 | // Re-export all core components for easier imports 10 | pub use dyn_tool::AnyOutput; 11 | pub use registry::ToolRegistry; 12 | pub use render::{Render, ResourcesTracker}; 13 | pub use result::ToolResult; 14 | pub use spec::{ToolScope, ToolSpec}; 15 | pub use tool::{Tool, ToolContext}; 16 | -------------------------------------------------------------------------------- /crates/code_assistant/src/tools/core/registry.rs: -------------------------------------------------------------------------------- 1 | use std::collections::HashMap; 2 | use std::sync::OnceLock; 3 | 4 | use crate::tools::core::dyn_tool::DynTool; 5 | use crate::tools::core::spec::ToolScope; 6 | use crate::tools::AnnotatedToolDefinition; 7 | 8 | /// Central registry for all tools in the system 9 | pub struct ToolRegistry { 10 | tools: HashMap>, 11 | } 12 | 13 | impl ToolRegistry { 14 | /// Get the global singleton instance of the registry 15 | pub fn global() -> &'static Self { 16 | // Singleton instance of the registry 17 | static INSTANCE: OnceLock = OnceLock::new(); 18 | INSTANCE.get_or_init(|| { 19 | let mut registry = ToolRegistry::new(); 20 | registry.register_default_tools(); 21 | registry 22 | }) 23 | } 24 | 25 | /// Create a new empty registry 26 | pub fn new() -> Self { 27 | Self { 28 | tools: HashMap::new(), 29 | } 30 | } 31 | 32 | /// Register a tool in the registry 33 | pub fn register(&mut self, tool: Box) { 34 | self.tools.insert(tool.spec().name.to_string(), tool); 35 | } 36 | 37 | /// Get a tool by name 38 | pub fn get(&self, name: &str) -> Option<&Box> { 39 | self.tools.get(name) 40 | } 41 | 42 | /// Get tool definitions for a specific mode 43 | pub fn get_tool_definitions_for_scope(&self, mode: ToolScope) -> Vec { 44 | self.tools 45 | .values() 46 | .filter(|tool| tool.spec().supported_scopes.contains(&mode)) 47 | .map(|tool| AnnotatedToolDefinition { 48 | name: tool.spec().name.to_string(), 49 | description: tool.spec().description.to_string(), 50 | parameters: tool.spec().parameters_schema.clone(), 51 | annotations: tool.spec().annotations.clone(), 52 | }) 53 | .collect() 54 | } 55 | 56 | /// Register all default tools in the system 57 | /// This will be expanded as we implement more tools 58 | fn register_default_tools(&mut self) { 59 | // Import all tools 60 | use crate::tools::impls::{ 61 | DeleteFilesTool, ExecuteCommandTool, ListFilesTool, ListProjectsTool, 62 | PerplexityAskTool, ReadFilesTool, ReplaceInFileTool, SearchFilesTool, WebFetchTool, 63 | WebSearchTool, WriteFileTool, 64 | }; 65 | 66 | // Register tools 67 | self.register(Box::new(DeleteFilesTool)); 68 | self.register(Box::new(ExecuteCommandTool)); 69 | self.register(Box::new(ListFilesTool)); 70 | self.register(Box::new(ListProjectsTool)); 71 | self.register(Box::new(PerplexityAskTool)); 72 | self.register(Box::new(ReadFilesTool)); 73 | self.register(Box::new(ReplaceInFileTool)); 74 | self.register(Box::new(SearchFilesTool)); 75 | self.register(Box::new(WebFetchTool)); 76 | self.register(Box::new(WebSearchTool)); 77 | self.register(Box::new(WriteFileTool)); 78 | 79 | // More tools will be added here as they are implemented 80 | } 81 | } 82 | -------------------------------------------------------------------------------- /crates/code_assistant/src/tools/core/render.rs: -------------------------------------------------------------------------------- 1 | use std::collections::HashSet; 2 | 3 | /// Responsible for formatting tool outputs for display 4 | pub trait Render: Send + Sync + 'static { 5 | /// Generate a short status message for display in action history 6 | fn status(&self) -> String; 7 | 8 | /// Format the detailed output, with awareness of other tool results 9 | /// The resources_tracker helps detect and handle redundant output 10 | fn render(&self, resources_tracker: &mut ResourcesTracker) -> String; 11 | } 12 | 13 | /// Tracks resources that have been included in tool outputs to prevent redundant display 14 | pub struct ResourcesTracker { 15 | /// Set of already rendered resource identifiers 16 | rendered_resources: HashSet, 17 | } 18 | 19 | impl ResourcesTracker { 20 | /// Create a new empty resources tracker 21 | pub fn new() -> Self { 22 | Self { 23 | rendered_resources: HashSet::new(), 24 | } 25 | } 26 | 27 | /// Check if a resource has already been rendered 28 | pub fn is_rendered(&self, resource_id: &str) -> bool { 29 | self.rendered_resources.contains(resource_id) 30 | } 31 | 32 | /// Mark a resource as rendered to prevent duplicate display 33 | pub fn mark_rendered(&mut self, resource_id: String) { 34 | self.rendered_resources.insert(resource_id); 35 | } 36 | } 37 | -------------------------------------------------------------------------------- /crates/code_assistant/src/tools/core/result.rs: -------------------------------------------------------------------------------- 1 | /// Trait for determining whether a tool execution was successful 2 | pub trait ToolResult: Send + Sync + 'static { 3 | /// Returns whether the tool execution was successful 4 | /// This is used for status reporting and can affect how the result is displayed 5 | fn is_success(&self) -> bool; 6 | } 7 | -------------------------------------------------------------------------------- /crates/code_assistant/src/tools/core/spec.rs: -------------------------------------------------------------------------------- 1 | /// Define available modes for tools 2 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] 3 | pub enum ToolScope { 4 | /// Tool can be used in the MCP server 5 | McpServer, 6 | /// Tool can be used in the message history agent 7 | Agent, 8 | } 9 | 10 | /// Specification for a tool, including metadata 11 | #[derive(Clone)] 12 | pub struct ToolSpec { 13 | /// Unique name of the tool 14 | pub name: &'static str, 15 | /// Detailed description of what the tool does 16 | pub description: &'static str, 17 | /// JSON Schema for the tool's parameters 18 | pub parameters_schema: serde_json::Value, 19 | /// Optional annotations for LLM-specific instructions 20 | pub annotations: Option, 21 | /// Which execution modes this tool supports 22 | pub supported_scopes: &'static [ToolScope], 23 | } 24 | -------------------------------------------------------------------------------- /crates/code_assistant/src/tools/core/tool.rs: -------------------------------------------------------------------------------- 1 | use super::render::Render; 2 | use super::result::ToolResult; 3 | use super::spec::ToolSpec; 4 | use crate::types::WorkingMemory; 5 | use anyhow::{anyhow, Result}; 6 | use serde::{de::DeserializeOwned, Deserialize, Serialize}; 7 | 8 | /// Context provided to tools during execution 9 | pub struct ToolContext<'a> { 10 | /// Project manager for accessing files 11 | pub project_manager: &'a dyn crate::config::ProjectManager, 12 | /// Command executor for running shell commands 13 | pub command_executor: &'a dyn crate::utils::CommandExecutor, 14 | /// Optional working memory (available in WorkingMemoryAgent mode) 15 | pub working_memory: Option<&'a mut WorkingMemory>, 16 | } 17 | 18 | /// Core trait for tools, defining the execution interface 19 | #[async_trait::async_trait] 20 | pub trait Tool: Send + Sync + 'static { 21 | /// Input type for this tool, must be deserializable from JSON 22 | type Input: DeserializeOwned + Send; 23 | 24 | /// Output type for this tool, must implement Render, ToolResult and Serialize/Deserialize 25 | type Output: Render + ToolResult + Serialize + for<'de> Deserialize<'de> + Send + Sync; 26 | 27 | /// Get the metadata for this tool 28 | fn spec(&self) -> ToolSpec; 29 | 30 | /// Execute the tool with the given context and input 31 | async fn execute<'a>( 32 | &self, 33 | context: &mut ToolContext<'a>, 34 | input: Self::Input, 35 | ) -> Result; 36 | 37 | /// Deserialize a JSON value into this tool's output type 38 | fn deserialize_output(&self, json: serde_json::Value) -> Result { 39 | serde_json::from_value(json).map_err(|e| anyhow!("Failed to deserialize output: {}", e)) 40 | } 41 | } 42 | -------------------------------------------------------------------------------- /crates/code_assistant/src/tools/impls/list_projects.rs: -------------------------------------------------------------------------------- 1 | use crate::tools::core::{ 2 | Render, ResourcesTracker, Tool, ToolContext, ToolResult, ToolScope, ToolSpec, 3 | }; 4 | use crate::types::Project; 5 | use anyhow::Result; 6 | use serde::{Deserialize, Serialize}; 7 | use serde_json::json; 8 | use std::collections::HashMap; 9 | 10 | // Input type (empty for this tool) 11 | #[derive(Deserialize)] 12 | pub struct ListProjectsInput {} 13 | 14 | // Output type 15 | #[derive(Serialize, Deserialize)] 16 | pub struct ListProjectsOutput { 17 | pub projects: HashMap, 18 | } 19 | 20 | // Render implementation for output formatting 21 | impl Render for ListProjectsOutput { 22 | fn status(&self) -> String { 23 | if self.projects.is_empty() { 24 | "No projects available".to_string() 25 | } else { 26 | format!("Found {} project(s)", self.projects.len()) 27 | } 28 | } 29 | 30 | fn render(&self, _tracker: &mut ResourcesTracker) -> String { 31 | if self.projects.is_empty() { 32 | return "No projects available".to_string(); 33 | } 34 | 35 | let mut output = String::from("Available projects:\n"); 36 | for (name, _) in &self.projects { 37 | output.push_str(&format!("- {}\n", name)); 38 | } 39 | 40 | output 41 | } 42 | } 43 | 44 | // ToolResult implementation 45 | impl ToolResult for ListProjectsOutput { 46 | fn is_success(&self) -> bool { 47 | true // Always successful even if no projects are found 48 | } 49 | } 50 | 51 | // The actual tool implementation 52 | pub struct ListProjectsTool; 53 | 54 | #[async_trait::async_trait] 55 | impl Tool for ListProjectsTool { 56 | type Input = ListProjectsInput; 57 | type Output = ListProjectsOutput; 58 | 59 | fn spec(&self) -> ToolSpec { 60 | let description = concat!( 61 | "List all available projects. ", 62 | "Use this tool to discover which projects are available for exploration." 63 | ); 64 | ToolSpec { 65 | name: "list_projects", 66 | description, 67 | parameters_schema: json!({ 68 | "type": "object", 69 | "properties": {}, 70 | "required": [] 71 | }), 72 | annotations: Some(json!({ 73 | "readOnlyHint": true 74 | })), 75 | supported_scopes: &[ToolScope::McpServer, ToolScope::Agent], 76 | } 77 | } 78 | 79 | async fn execute<'a>( 80 | &self, 81 | context: &mut ToolContext<'a>, 82 | _input: Self::Input, 83 | ) -> Result { 84 | // Load projects using the ProjectManager from the context 85 | let projects = context.project_manager.get_projects()?; 86 | 87 | Ok(ListProjectsOutput { projects }) 88 | } 89 | } 90 | 91 | #[cfg(test)] 92 | mod tests { 93 | use super::*; 94 | 95 | #[tokio::test] 96 | async fn test_list_projects_renders_correctly() { 97 | // Create sample data 98 | let mut projects = HashMap::new(); 99 | projects.insert( 100 | "test-project".to_string(), 101 | Project { 102 | path: std::path::PathBuf::from("/path/to/test-project"), 103 | }, 104 | ); 105 | 106 | let output = ListProjectsOutput { projects }; 107 | let mut tracker = ResourcesTracker::new(); 108 | 109 | // Test rendering 110 | let rendered = output.render(&mut tracker); 111 | assert!(rendered.contains("Available projects:")); 112 | assert!(rendered.contains("- test-project")); 113 | } 114 | } 115 | -------------------------------------------------------------------------------- /crates/code_assistant/src/tools/impls/mod.rs: -------------------------------------------------------------------------------- 1 | // Tool implementations 2 | pub mod delete_files; 3 | pub mod execute_command; 4 | pub mod list_files; 5 | pub mod list_projects; 6 | pub mod perplexity_ask; 7 | pub mod read_files; 8 | pub mod replace_in_file; 9 | pub mod search_files; 10 | pub mod web_fetch; 11 | pub mod web_search; 12 | pub mod write_file; 13 | 14 | // Re-export all tools for registration 15 | pub use delete_files::DeleteFilesTool; 16 | pub use execute_command::ExecuteCommandTool; 17 | pub use list_files::ListFilesTool; 18 | pub use list_projects::ListProjectsTool; 19 | pub use perplexity_ask::PerplexityAskTool; 20 | pub use read_files::ReadFilesTool; 21 | pub use replace_in_file::ReplaceInFileTool; 22 | pub use search_files::SearchFilesTool; 23 | pub use web_fetch::WebFetchTool; 24 | pub use web_search::WebSearchTool; 25 | pub use write_file::WriteFileTool; 26 | -------------------------------------------------------------------------------- /crates/code_assistant/src/tools/mod.rs: -------------------------------------------------------------------------------- 1 | // Original tools implementation 2 | mod parse; 3 | mod types; 4 | 5 | // New trait-based tools implementation 6 | pub mod core; 7 | pub mod impls; 8 | 9 | #[cfg(test)] 10 | mod tests; 11 | 12 | pub use parse::{parse_tool_xml, TOOL_TAG_PREFIX}; 13 | pub use types::AnnotatedToolDefinition; 14 | -------------------------------------------------------------------------------- /crates/code_assistant/src/tools/types.rs: -------------------------------------------------------------------------------- 1 | use llm::ToolDefinition; 2 | use serde::{Deserialize, Serialize}; 3 | 4 | /// Enhanced version of the base ToolDefinition with additional metadata fields 5 | #[derive(Debug, Clone, Serialize, Deserialize)] 6 | pub struct AnnotatedToolDefinition { 7 | pub name: String, 8 | pub description: String, 9 | pub parameters: serde_json::Value, 10 | #[serde(skip_serializing_if = "Option::is_none")] 11 | pub annotations: Option, 12 | } 13 | 14 | impl AnnotatedToolDefinition { 15 | /// Convert to a basic ToolDefinition (without annotations) for LLM providers 16 | pub fn to_tool_definition(&self) -> ToolDefinition { 17 | ToolDefinition { 18 | name: self.name.clone(), 19 | description: self.description.clone(), 20 | parameters: self.parameters.clone(), 21 | } 22 | } 23 | 24 | /// Convert a vector of AnnotatedToolDefinition to a vector of ToolDefinition 25 | pub fn to_tool_definitions(tools: Vec) -> Vec { 26 | tools.into_iter().map(|t| t.to_tool_definition()).collect() 27 | } 28 | } 29 | -------------------------------------------------------------------------------- /crates/code_assistant/src/ui/gpui/assets.rs: -------------------------------------------------------------------------------- 1 | use anyhow::anyhow; 2 | use gpui::{AssetSource, Result, SharedString}; 3 | use rust_embed::RustEmbed; 4 | 5 | #[derive(RustEmbed)] 6 | #[folder = "assets"] 7 | #[include = "icons/**/*"] 8 | #[exclude = "*.DS_Store"] 9 | pub struct Assets; 10 | 11 | impl AssetSource for Assets { 12 | fn load(&self, path: &str) -> Result>> { 13 | Self::get(path) 14 | .map(|f| Some(f.data)) 15 | .ok_or_else(|| anyhow!("could not find asset at path \"{}\"", path)) 16 | } 17 | 18 | fn list(&self, path: &str) -> Result> { 19 | Ok(Self::iter() 20 | .filter_map(|p| { 21 | if p.starts_with(path) { 22 | Some(p.into()) 23 | } else { 24 | None 25 | } 26 | }) 27 | .collect()) 28 | } 29 | } 30 | -------------------------------------------------------------------------------- /crates/code_assistant/src/ui/gpui/content_renderer.rs: -------------------------------------------------------------------------------- 1 | use crate::ui::gpui::parameter_renderers::ParameterRenderer; 2 | use gpui::{div, px, rgba, Element, FontWeight, ParentElement, Styled}; 3 | 4 | /// Renderer for the "content" parameter of the "write_file" tool 5 | pub struct ContentRenderer; 6 | 7 | impl ParameterRenderer for ContentRenderer { 8 | fn supported_parameters(&self) -> Vec<(String, String)> { 9 | vec![("write_file".to_string(), "content".to_string())] 10 | } 11 | 12 | fn render( 13 | &self, 14 | _tool_name: &str, 15 | _param_name: &str, 16 | param_value: &str, 17 | theme: &gpui_component::theme::Theme, 18 | ) -> gpui::AnyElement { 19 | // Container for the content - no parameter name shown 20 | div() 21 | .rounded_md() 22 | .bg(if theme.is_dark() { 23 | rgba(0x0A0A0AFF) // Darker background in Dark Mode 24 | } else { 25 | rgba(0xEAEAEAFF) // Lighter background in Light Mode 26 | }) 27 | .p_2() 28 | .text_size(px(15.)) 29 | .font_weight(FontWeight(500.0)) 30 | .child(param_value.to_string()) 31 | .into_any() 32 | } 33 | 34 | fn is_full_width(&self, _tool_name: &str, _param_name: &str) -> bool { 35 | true // Content parameter is always full-width 36 | } 37 | } 38 | -------------------------------------------------------------------------------- /crates/code_assistant/src/ui/gpui/parameter_renderers.rs: -------------------------------------------------------------------------------- 1 | use gpui::{px, Element, IntoElement, ParentElement, Styled}; 2 | use std::collections::HashMap; 3 | use std::sync::{Arc, Mutex, OnceLock}; 4 | use tracing::warn; 5 | 6 | /// A unique key for tool+parameter combinations 7 | pub type ParameterKey = String; 8 | 9 | /// Helper function to create a unique key for a tool-parameter combination 10 | pub fn create_parameter_key(tool_name: &str, param_name: &str) -> ParameterKey { 11 | format!("{}:{}", tool_name, param_name) 12 | } 13 | 14 | /// Trait for parameter renderers that can provide custom rendering for tool parameters 15 | pub trait ParameterRenderer: Send + Sync { 16 | /// List of supported tool+parameter combinations 17 | fn supported_parameters(&self) -> Vec<(String, String)>; 18 | 19 | /// Render the parameter as a UI element 20 | fn render( 21 | &self, 22 | tool_name: &str, 23 | param_name: &str, 24 | param_value: &str, 25 | theme: &gpui_component::theme::Theme, 26 | ) -> gpui::AnyElement; 27 | 28 | /// Indicates if this parameter should be rendered with full width 29 | /// Default is false (normal inline parameter) 30 | fn is_full_width(&self, _tool_name: &str, _param_name: &str) -> bool { 31 | false 32 | } 33 | } 34 | 35 | /// Registry for parameter renderers 36 | pub struct ParameterRendererRegistry { 37 | // Map from tool+parameter key to renderer 38 | renderers: HashMap>>, 39 | // Default renderer for parameters with no specific renderer 40 | default_renderer: Arc>, 41 | } 42 | 43 | // Global registry singleton using OnceLock (thread-safe) 44 | static GLOBAL_REGISTRY: OnceLock>>> = OnceLock::new(); 45 | 46 | impl ParameterRendererRegistry { 47 | // Set the global registry 48 | pub fn set_global(registry: Arc) { 49 | // Initialize the global mutex if not already initialized 50 | let global_mutex = GLOBAL_REGISTRY.get_or_init(|| Mutex::new(None)); 51 | 52 | // Set the registry instance 53 | if let Ok(mut guard) = global_mutex.lock() { 54 | *guard = Some(registry); 55 | } else { 56 | warn!("Failed to acquire lock for setting global registry"); 57 | } 58 | } 59 | 60 | // Get a reference to the global registry 61 | pub fn global() -> Option> { 62 | if let Some(global_mutex) = GLOBAL_REGISTRY.get() { 63 | if let Ok(guard) = global_mutex.lock() { 64 | return guard.clone(); 65 | } 66 | } 67 | None 68 | } 69 | 70 | /// Create a new registry with the given default renderer 71 | pub fn new(default_renderer: Box) -> Self { 72 | Self { 73 | renderers: HashMap::new(), 74 | default_renderer: Arc::new(default_renderer), 75 | } 76 | } 77 | 78 | /// Register a new renderer for its supported parameters 79 | pub fn register_renderer(&mut self, renderer: Box) { 80 | let renderer_arc = Arc::new(renderer); 81 | 82 | for (tool_name, param_name) in renderer_arc.supported_parameters() { 83 | let key = create_parameter_key(&tool_name, ¶m_name); 84 | if self.renderers.contains_key(&key) { 85 | warn!("Overriding existing renderer for {}", key); 86 | } 87 | self.renderers.insert(key, renderer_arc.clone()); 88 | } 89 | } 90 | 91 | /// Get the appropriate renderer for a tool+parameter combination 92 | pub fn get_renderer( 93 | &self, 94 | tool_name: &str, 95 | param_name: &str, 96 | ) -> Arc> { 97 | let key = create_parameter_key(tool_name, param_name); 98 | 99 | self.renderers 100 | .get(&key) 101 | .unwrap_or(&self.default_renderer) 102 | .clone() 103 | } 104 | 105 | /// Render a parameter using the appropriate renderer 106 | pub fn render_parameter( 107 | &self, 108 | tool_name: &str, 109 | param_name: &str, 110 | param_value: &str, 111 | theme: &gpui_component::theme::Theme, 112 | ) -> gpui::AnyElement { 113 | let renderer = self.get_renderer(tool_name, param_name); 114 | renderer.render(tool_name, param_name, param_value, theme) 115 | } 116 | } 117 | 118 | /// Default parameter renderer that displays parameters in a simple badge format 119 | pub struct DefaultParameterRenderer; 120 | 121 | impl ParameterRenderer for DefaultParameterRenderer { 122 | fn supported_parameters(&self) -> Vec<(String, String)> { 123 | // Default renderer supports no specific parameters 124 | Vec::new() 125 | } 126 | 127 | fn render( 128 | &self, 129 | _tool_name: &str, 130 | param_name: &str, 131 | param_value: &str, 132 | theme: &gpui_component::theme::Theme, 133 | ) -> gpui::AnyElement { 134 | use gpui::{div, FontWeight}; 135 | 136 | div() 137 | .rounded_md() 138 | .px_2() 139 | .py_1() 140 | .text_size(px(13.)) 141 | .bg(crate::ui::gpui::theme::colors::tool_parameter_bg(theme)) 142 | .child( 143 | div() 144 | .flex() 145 | .flex_row() 146 | .items_center() 147 | .gap_1() 148 | .children(vec![ 149 | div() 150 | .font_weight(FontWeight(500.0)) 151 | .text_color(crate::ui::gpui::theme::colors::tool_parameter_label(theme)) 152 | .child(format!("{}:", param_name)) 153 | .into_any(), 154 | div() 155 | .text_color(crate::ui::gpui::theme::colors::tool_parameter_value(theme)) 156 | .child(param_value.to_string()) 157 | .into_any(), 158 | ]), 159 | ) 160 | .into_any_element() 161 | } 162 | } 163 | -------------------------------------------------------------------------------- /crates/code_assistant/src/ui/gpui/path_util.rs: -------------------------------------------------------------------------------- 1 | use std::path::Path; 2 | 3 | /// Helper trait to extend Path with icon-related functionality. 4 | pub trait PathExt { 5 | /// Returns either the suffix if available, or the file stem otherwise to determine 6 | /// which file icon to use. 7 | fn icon_stem_or_suffix(&self) -> Option<&str>; 8 | } 9 | 10 | impl> PathExt for T { 11 | fn icon_stem_or_suffix(&self) -> Option<&str> { 12 | let path = self.as_ref(); 13 | let file_name = path.file_name()?.to_str()?; 14 | 15 | // For hidden files (Unix style), return the name without the leading dot 16 | if file_name.starts_with('.') { 17 | return file_name.strip_prefix('.'); 18 | } 19 | 20 | // Try to get extension, or fall back to file stem 21 | path.extension() 22 | .and_then(|e| e.to_str()) 23 | .or_else(|| path.file_stem()?.to_str()) 24 | } 25 | } 26 | 27 | #[cfg(test)] 28 | mod tests { 29 | use super::*; 30 | 31 | #[test] 32 | fn test_icon_stem_or_suffix() { 33 | // No dots in name 34 | let path = Path::new("/a/b/c/file_name.rs"); 35 | assert_eq!(path.icon_stem_or_suffix(), Some("rs")); 36 | 37 | // Single dot in name 38 | let path = Path::new("/a/b/c/file.name.rs"); 39 | assert_eq!(path.icon_stem_or_suffix(), Some("rs")); 40 | 41 | // No suffix 42 | let path = Path::new("/a/b/c/file"); 43 | assert_eq!(path.icon_stem_or_suffix(), Some("file")); 44 | 45 | // Multiple dots in name 46 | let path = Path::new("/a/b/c/long.file.name.rs"); 47 | assert_eq!(path.icon_stem_or_suffix(), Some("rs")); 48 | 49 | // Hidden file, no extension 50 | let path = Path::new("/a/b/c/.gitignore"); 51 | assert_eq!(path.icon_stem_or_suffix(), Some("gitignore")); 52 | 53 | // Hidden file, with extension 54 | let path = Path::new("/a/b/c/.eslintrc.js"); 55 | assert_eq!(path.icon_stem_or_suffix(), Some("eslintrc.js")); 56 | } 57 | } 58 | -------------------------------------------------------------------------------- /crates/code_assistant/src/ui/gpui/simple_renderers.rs: -------------------------------------------------------------------------------- 1 | use crate::ui::gpui::parameter_renderers::ParameterRenderer; 2 | use gpui::{div, px, IntoElement, ParentElement, Styled}; 3 | 4 | /// Renderer for parameters that shouldn't show their parameter name 5 | pub struct SimpleParameterRenderer { 6 | /// The list of tool+parameter combinations that should use this renderer 7 | supported_combinations: Vec<(String, String)>, 8 | /// Whether this parameter should be rendered with full width 9 | full_width: bool, 10 | } 11 | 12 | impl SimpleParameterRenderer { 13 | /// Create a new simple parameter renderer with specified combinations 14 | pub fn new(combinations: Vec<(String, String)>, full_width: bool) -> Self { 15 | Self { 16 | supported_combinations: combinations, 17 | full_width, 18 | } 19 | } 20 | } 21 | 22 | impl ParameterRenderer for SimpleParameterRenderer { 23 | fn supported_parameters(&self) -> Vec<(String, String)> { 24 | self.supported_combinations.clone() 25 | } 26 | 27 | fn render( 28 | &self, 29 | _tool_name: &str, 30 | _param_name: &str, 31 | param_value: &str, 32 | theme: &gpui_component::theme::Theme, 33 | ) -> gpui::AnyElement { 34 | div() 35 | .rounded_md() 36 | .px_2() 37 | .py_1() 38 | .text_size(px(13.)) 39 | .bg(crate::ui::gpui::theme::colors::tool_parameter_bg(theme)) 40 | .text_color(crate::ui::gpui::theme::colors::tool_parameter_value(theme)) 41 | .child(param_value.to_string()) 42 | .into_any_element() 43 | } 44 | 45 | fn is_full_width(&self, _tool_name: &str, _param_name: &str) -> bool { 46 | self.full_width 47 | } 48 | } 49 | -------------------------------------------------------------------------------- /crates/code_assistant/src/ui/gpui/ui_events.rs: -------------------------------------------------------------------------------- 1 | use crate::types::WorkingMemory; 2 | use crate::ui::gpui::elements::MessageRole; 3 | use crate::ui::ToolStatus; 4 | 5 | /// Events for UI updates from the agent thread 6 | #[derive(Debug, Clone)] 7 | pub enum UiEvent { 8 | /// Display a new message or append to an existing one 9 | DisplayMessage { content: String, role: MessageRole }, 10 | /// Append to the last text block 11 | AppendToTextBlock { content: String }, 12 | /// Append to the last thinking block 13 | AppendToThinkingBlock { content: String }, 14 | /// Start a tool invocation 15 | StartTool { name: String, id: String }, 16 | /// Add or update a tool parameter 17 | UpdateToolParameter { 18 | tool_id: String, 19 | name: String, 20 | value: String, 21 | }, 22 | /// Update a tool status 23 | UpdateToolStatus { 24 | tool_id: String, 25 | status: ToolStatus, 26 | message: Option, 27 | output: Option, 28 | }, 29 | /// End a tool invocation 30 | EndTool { id: String }, 31 | /// Update the working memory view 32 | UpdateMemory { memory: WorkingMemory }, 33 | /// Streaming started for a request 34 | StreamingStarted(u64), 35 | /// Streaming stopped for a request 36 | StreamingStopped { id: u64, cancelled: bool }, 37 | } 38 | -------------------------------------------------------------------------------- /crates/code_assistant/src/ui/mod.rs: -------------------------------------------------------------------------------- 1 | pub mod gpui; 2 | pub mod streaming; 3 | pub mod terminal; 4 | use crate::types::WorkingMemory; 5 | use async_trait::async_trait; 6 | pub use streaming::DisplayFragment; 7 | use thiserror::Error; 8 | 9 | #[derive(Debug, Clone, Copy, PartialEq)] 10 | pub enum ToolStatus { 11 | Pending, // Default status when a tool appears in the stream 12 | Running, // Tool is currently being executed 13 | Success, // Execution was successful 14 | Error, // Error during execution 15 | } 16 | 17 | #[derive(Debug, Clone, Copy, PartialEq)] 18 | pub enum StreamingState { 19 | Idle, // No active streaming, ready to send 20 | Streaming, // Currently streaming response 21 | StopRequested, // User requested stop, waiting for stream to end 22 | } 23 | 24 | #[derive(Debug, Clone)] 25 | pub enum UIMessage { 26 | // System actions that the agent takes 27 | Action(String), 28 | // User input messages 29 | UserInput(String), 30 | } 31 | 32 | #[derive(Error, Debug)] 33 | pub enum UIError { 34 | #[error("IO error: {0}")] 35 | IOError(#[from] std::io::Error), 36 | // #[error("Input cancelled")] 37 | // Cancelled, 38 | // #[error("Other UI error: {0}")] 39 | // Other(String), 40 | } 41 | 42 | #[async_trait] 43 | pub trait UserInterface: Send + Sync { 44 | /// Display a message to the user 45 | async fn display(&self, message: UIMessage) -> Result<(), UIError>; 46 | 47 | /// Get input from the user 48 | async fn get_input(&self) -> Result; 49 | 50 | /// Display a streaming fragment with specific type information 51 | fn display_fragment(&self, fragment: &DisplayFragment) -> Result<(), UIError>; 52 | 53 | /// Update tool status for a specific tool 54 | async fn update_tool_status( 55 | &self, 56 | tool_id: &str, 57 | status: ToolStatus, 58 | message: Option, 59 | output: Option, 60 | ) -> Result<(), UIError>; 61 | 62 | /// Update memory view with current working memory 63 | async fn update_memory(&self, memory: &WorkingMemory) -> Result<(), UIError>; 64 | 65 | /// Informs the UI that a new LLM request is starting 66 | /// Returns the request ID that can be used to correlate tool invocations 67 | async fn begin_llm_request(&self) -> Result; 68 | 69 | /// Informs the UI that an LLM request has completed 70 | async fn end_llm_request(&self, request_id: u64, cancelled: bool) -> Result<(), UIError>; 71 | 72 | /// Check if streaming should continue 73 | fn should_streaming_continue(&self) -> bool; 74 | } 75 | 76 | #[cfg(test)] 77 | mod terminal_test; 78 | -------------------------------------------------------------------------------- /crates/code_assistant/src/ui/streaming/mod.rs: -------------------------------------------------------------------------------- 1 | //! Streaming processor for handling chunks from LLM providers 2 | 3 | use crate::agent::ToolMode; 4 | use crate::ui::UIError; 5 | use crate::ui::UserInterface; 6 | use llm::StreamingChunk; 7 | use std::sync::Arc; 8 | 9 | mod json_processor; 10 | mod xml_processor; 11 | 12 | #[cfg(test)] 13 | mod json_processor_tests; 14 | #[cfg(test)] 15 | mod test_utils; 16 | #[cfg(test)] 17 | mod xml_processor_tests; 18 | 19 | /// Fragments for display in UI components 20 | #[derive(Debug, Clone, PartialEq)] 21 | pub enum DisplayFragment { 22 | /// Regular plain text 23 | PlainText(String), 24 | /// Thinking text (shown differently) 25 | ThinkingText(String), 26 | /// Tool invocation start 27 | ToolName { name: String, id: String }, 28 | /// Parameter for a tool 29 | ToolParameter { 30 | name: String, 31 | value: String, 32 | tool_id: String, 33 | }, 34 | /// End of a tool invocation 35 | ToolEnd { id: String }, 36 | } 37 | 38 | /// Common trait for stream processors 39 | pub trait StreamProcessorTrait: Send + Sync { 40 | /// Create a new stream processor with the given UI 41 | fn new(ui: Arc>) -> Self 42 | where 43 | Self: Sized; 44 | 45 | /// Process a streaming chunk and send display fragments to the UI 46 | fn process(&mut self, chunk: &StreamingChunk) -> Result<(), UIError>; 47 | } 48 | 49 | // Export the concrete implementations 50 | pub use json_processor::JsonStreamProcessor; 51 | pub use xml_processor::XmlStreamProcessor; 52 | 53 | /// Factory function to create the appropriate processor based on tool mode 54 | pub fn create_stream_processor( 55 | tool_mode: ToolMode, 56 | ui: Arc>, 57 | ) -> Box { 58 | match tool_mode { 59 | ToolMode::Xml => Box::new(XmlStreamProcessor::new(ui)), 60 | ToolMode::Native => Box::new(JsonStreamProcessor::new(ui)), 61 | } 62 | } 63 | -------------------------------------------------------------------------------- /crates/code_assistant/src/ui/terminal_test.rs: -------------------------------------------------------------------------------- 1 | //! Tests for the terminal UI formatting and output 2 | 3 | use super::streaming::DisplayFragment; 4 | use super::terminal::TerminalUI; 5 | use super::UserInterface; 6 | use std::io::Write; 7 | use std::sync::{Arc, Mutex}; 8 | 9 | // Mock stdout to capture output 10 | struct TestWriter { 11 | buffer: Vec, 12 | } 13 | 14 | impl TestWriter { 15 | fn new() -> Self { 16 | Self { buffer: Vec::new() } 17 | } 18 | 19 | fn get_output(&self) -> String { 20 | String::from_utf8_lossy(&self.buffer).to_string() 21 | } 22 | } 23 | 24 | impl Write for TestWriter { 25 | fn write(&mut self, buf: &[u8]) -> std::io::Result { 26 | self.buffer.extend_from_slice(buf); 27 | Ok(buf.len()) 28 | } 29 | 30 | fn flush(&mut self) -> std::io::Result<()> { 31 | Ok(()) 32 | } 33 | } 34 | 35 | // Helper function to create a terminal UI with a test writer 36 | fn create_test_terminal_ui() -> (TerminalUI, Arc>) { 37 | let writer = Arc::new(Mutex::new(TestWriter::new())); 38 | 39 | // Create a wrapper to satisfy the trait bounds 40 | struct WriterWrapper(Arc>); 41 | 42 | impl Write for WriterWrapper { 43 | fn write(&mut self, buf: &[u8]) -> std::io::Result { 44 | self.0.lock().unwrap().write(buf) 45 | } 46 | 47 | fn flush(&mut self) -> std::io::Result<()> { 48 | self.0.lock().unwrap().flush() 49 | } 50 | } 51 | 52 | let wrapper = Box::new(WriterWrapper(writer.clone())); 53 | let ui = TerminalUI::with_test_writer(wrapper); 54 | 55 | (ui, writer) 56 | } 57 | 58 | #[test] 59 | fn test_terminal_formatting() { 60 | // Create terminal UI with test writer 61 | let (ui, writer) = create_test_terminal_ui(); 62 | 63 | // Test various display fragments 64 | ui.display_fragment(&DisplayFragment::PlainText("Hello world".to_string())) 65 | .unwrap(); 66 | ui.display_fragment(&DisplayFragment::ThinkingText("Thinking...".to_string())) 67 | .unwrap(); 68 | 69 | // The newline is needed because the tool name formatting starts with a newline 70 | ui.display_fragment(&DisplayFragment::PlainText("\n".to_string())) 71 | .unwrap(); 72 | 73 | ui.display_fragment(&DisplayFragment::ToolName { 74 | name: "search_files".to_string(), 75 | id: "tool-123".to_string(), 76 | }) 77 | .unwrap(); 78 | ui.display_fragment(&DisplayFragment::ToolParameter { 79 | name: "query".to_string(), 80 | value: "search term".to_string(), 81 | tool_id: "tool-123".to_string(), 82 | }) 83 | .unwrap(); 84 | 85 | // Check the output 86 | let output = writer.lock().unwrap().get_output(); 87 | println!("Output:\n{}", output); 88 | 89 | // Verify various formatting aspects 90 | assert!( 91 | output.contains("Hello world"), 92 | "Plain text should be displayed as-is" 93 | ); 94 | 95 | // Thinking text should be styled (usually italic and grey, but we can't easily check styling in tests) 96 | assert!( 97 | output.contains("Thinking..."), 98 | "Thinking text should be visible" 99 | ); 100 | 101 | // We can check if the bullet point appears 102 | assert!(output.contains("•"), "Bullet point should be visible"); 103 | 104 | // Verify tool name appears 105 | assert!( 106 | output.contains("search_files"), 107 | "Tool name should be visible" 108 | ); 109 | 110 | // Parameter should be formatted with indentation and name 111 | // In der Ausgabe mit ANSI-Farbcodes ist es schwer, genau nach "query:" zu suchen 112 | // Wir prüfen stattdessen, ob der Wert vorhanden ist 113 | assert!( 114 | output.contains("search term"), 115 | "Parameter value should be visible" 116 | ); 117 | } 118 | -------------------------------------------------------------------------------- /crates/code_assistant/src/utils/command.rs: -------------------------------------------------------------------------------- 1 | use anyhow::Result; 2 | use std::path::PathBuf; 3 | 4 | #[derive(Clone)] 5 | pub struct CommandOutput { 6 | pub success: bool, 7 | pub output: String, 8 | } 9 | 10 | #[async_trait::async_trait] 11 | pub trait CommandExecutor: Send + Sync { 12 | async fn execute( 13 | &self, 14 | command_line: &str, 15 | working_dir: Option<&PathBuf>, 16 | ) -> Result; 17 | } 18 | 19 | pub struct DefaultCommandExecutor; 20 | 21 | #[async_trait::async_trait] 22 | impl CommandExecutor for DefaultCommandExecutor { 23 | async fn execute( 24 | &self, 25 | command_line: &str, 26 | working_dir: Option<&PathBuf>, 27 | ) -> Result { 28 | // Validate working_dir first 29 | if let Some(dir) = working_dir { 30 | if !dir.exists() { 31 | return Err(anyhow::anyhow!( 32 | "Working directory does not exist: {}", 33 | dir.display() 34 | )); 35 | } 36 | if !dir.is_dir() { 37 | return Err(anyhow::anyhow!( 38 | "Path is not a directory: {}", 39 | dir.display() 40 | )); 41 | } 42 | } 43 | 44 | // Create shell command using login shell or fallback 45 | #[cfg(target_family = "unix")] 46 | let shell = std::env::var("SHELL").unwrap_or_else(|_| "/bin/bash".to_string()); 47 | #[cfg(target_family = "unix")] 48 | let mut cmd = std::process::Command::new(shell); 49 | #[cfg(target_family = "unix")] 50 | cmd.args(["-c", &format!("{} 2>&1", command_line)]); 51 | 52 | #[cfg(target_family = "windows")] 53 | let mut cmd = std::process::Command::new("cmd"); 54 | #[cfg(target_family = "windows")] 55 | cmd.args(["/C", &format!("{} 2>&1", command_line)]); 56 | 57 | if let Some(dir) = working_dir { 58 | cmd.current_dir(dir); 59 | } 60 | let output = cmd.output()?; 61 | 62 | Ok(CommandOutput { 63 | success: output.status.success(), 64 | output: String::from_utf8_lossy(&output.stdout).into_owned(), 65 | }) 66 | } 67 | } 68 | -------------------------------------------------------------------------------- /crates/code_assistant/src/utils/file_updater.rs: -------------------------------------------------------------------------------- 1 | use crate::types::FileReplacement; 2 | use crate::utils::encoding; 3 | 4 | #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] 5 | pub enum FileUpdaterError { 6 | SearchBlockNotFound(usize, String), 7 | MultipleMatches(usize, usize, String), 8 | Other(String), 9 | } 10 | 11 | impl std::fmt::Display for FileUpdaterError { 12 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 13 | match self { 14 | FileUpdaterError::SearchBlockNotFound(index, ..) => { 15 | write!( 16 | f, 17 | "Could not find SEARCH block with index {} in the file contents", 18 | index 19 | ) 20 | } 21 | FileUpdaterError::MultipleMatches(count, index, _) => { 22 | write!(f, "Found {} occurrences of SEARCH block with index {}\nA SEARCH block must match exactly one location. Try enlarging the section to replace.", count, index) 23 | } 24 | FileUpdaterError::Other(msg) => { 25 | write!(f, "{}", msg) 26 | } 27 | } 28 | } 29 | } 30 | 31 | impl std::error::Error for FileUpdaterError {} 32 | 33 | /// Apply replacements with content normalization to make SEARCH blocks more robust 34 | /// against whitespace and line ending differences 35 | pub fn apply_replacements_normalized( 36 | content: &str, 37 | replacements: &[FileReplacement], 38 | ) -> Result { 39 | // Normalize the input content first 40 | let normalized_content = encoding::normalize_content(content); 41 | let mut result = normalized_content.clone(); 42 | 43 | for (index, replacement) in replacements.iter().enumerate() { 44 | // Normalize the search string as well 45 | let normalized_search = encoding::normalize_content(&replacement.search); 46 | 47 | // Count occurrences 48 | let matches: Vec<_> = result.match_indices(&normalized_search).collect(); 49 | 50 | if matches.is_empty() { 51 | return Err( 52 | FileUpdaterError::SearchBlockNotFound(index, replacement.search.clone()).into(), 53 | ); 54 | } 55 | 56 | // Normalize the replace string as well 57 | let normalized_replace = encoding::normalize_content(&replacement.replace); 58 | 59 | if replacement.replace_all { 60 | // Replace all occurrences 61 | result = result.replace(&normalized_search, &normalized_replace); 62 | } else { 63 | // Exact-match mode: must have exactly one occurrence 64 | if matches.len() > 1 { 65 | return Err(FileUpdaterError::MultipleMatches( 66 | matches.len(), 67 | index, 68 | replacement.search.clone(), 69 | ) 70 | .into()); 71 | } 72 | 73 | // Replace the single occurrence 74 | let (pos, _) = matches[0]; 75 | result.replace_range(pos..pos + normalized_search.len(), &normalized_replace); 76 | } 77 | } 78 | 79 | Ok(result) 80 | } 81 | 82 | #[test] 83 | fn test_apply_replacements_normalized() -> Result<(), anyhow::Error> { 84 | let test_cases: Vec<(&str, Vec, Result<&str, &str>)> = vec![ 85 | // Test with trailing whitespace 86 | ( 87 | "Hello World \nThis is a test\nGoodbye", 88 | vec![FileReplacement { 89 | search: "Hello World\nThis".to_string(), // No trailing space in search 90 | replace: "Hi there\nNew".to_string(), 91 | replace_all: false, 92 | }], 93 | Ok("Hi there\nNew is a test\nGoodbye"), 94 | ), 95 | // Test with different line endings 96 | ( 97 | "function test() {\r\n console.log('test');\r\n}", // CRLF endings 98 | vec![FileReplacement { 99 | search: "function test() {\n console.log('test');\n}".to_string(), // LF endings 100 | replace: "function answer() {\n return 42;\n}".to_string(), 101 | replace_all: false, 102 | }], 103 | Ok("function answer() {\n return 42;\n}"), 104 | ), 105 | // Test with both line ending and whitespace differences 106 | ( 107 | "test line \r\nwith trailing space \r\nand CRLF endings", 108 | vec![FileReplacement { 109 | search: "test line\nwith trailing space\nand CRLF endings".to_string(), 110 | replace: "replaced content".to_string(), 111 | replace_all: false, 112 | }], 113 | Ok("replaced content"), 114 | ), 115 | // Test replacing all occurrences 116 | ( 117 | "log('test');\nlog('test2');\nlog('test3');", 118 | vec![FileReplacement { 119 | search: "log(".to_string(), 120 | replace: "console.log(".to_string(), 121 | replace_all: true, 122 | }], 123 | Ok("console.log('test');\nconsole.log('test2');\nconsole.log('test3');"), 124 | ), 125 | // Test error when multiple matches but replace_all is false 126 | ( 127 | "log('test');\nlog('test2');\nlog('test3');", 128 | vec![FileReplacement { 129 | search: "log(".to_string(), 130 | replace: "console.log(".to_string(), 131 | replace_all: false, 132 | }], 133 | Err("Found 3 occurrences"), 134 | ), 135 | ]; 136 | 137 | for (input, replacements, expected) in test_cases { 138 | let result = apply_replacements_normalized(input, &replacements); 139 | match (&result, &expected) { 140 | (Ok(res), Ok(exp)) => assert_eq!(res, exp), 141 | (Err(e), Err(exp)) => assert!(e.to_string().contains(exp)), 142 | _ => { 143 | panic!( 144 | "Test case result did not match expected outcome:\nResult: {:?}\nExpected: {:?}", 145 | result, expected 146 | ); 147 | } 148 | } 149 | } 150 | 151 | Ok(()) 152 | } 153 | -------------------------------------------------------------------------------- /crates/code_assistant/src/utils/mod.rs: -------------------------------------------------------------------------------- 1 | mod command; 2 | mod file_updater; 3 | mod writer; 4 | 5 | pub mod encoding; 6 | 7 | #[allow(unused_imports)] 8 | pub use command::{CommandExecutor, CommandOutput, DefaultCommandExecutor}; 9 | pub use file_updater::{apply_replacements_normalized, FileUpdaterError}; 10 | #[cfg(test)] 11 | pub use writer::MockWriter; 12 | pub use writer::{MessageWriter, StdoutWriter}; 13 | -------------------------------------------------------------------------------- /crates/code_assistant/src/utils/writer.rs: -------------------------------------------------------------------------------- 1 | use anyhow::Result; 2 | use async_trait::async_trait; 3 | use tokio::io::{AsyncWriteExt, Stdout}; 4 | 5 | #[cfg(test)] 6 | use std::sync::Arc; 7 | #[cfg(test)] 8 | use tokio::sync::Mutex as TokioMutex; 9 | 10 | /// A trait for writing messages to an output stream. 11 | /// This abstraction allows replacing the actual output writer in tests. 12 | #[async_trait] 13 | pub trait MessageWriter: Send + Sync { 14 | /// Write a message to the output stream and flush it. 15 | async fn write_message(&mut self, message: &str) -> Result<()>; 16 | } 17 | 18 | /// The default implementation of MessageWriter that writes to Stdout. 19 | pub struct StdoutWriter { 20 | stdout: Stdout, 21 | } 22 | 23 | impl StdoutWriter { 24 | pub fn new(stdout: Stdout) -> Self { 25 | Self { stdout } 26 | } 27 | } 28 | 29 | #[async_trait] 30 | impl MessageWriter for StdoutWriter { 31 | async fn write_message(&mut self, message: &str) -> Result<()> { 32 | self.stdout.write_all(message.as_bytes()).await?; 33 | self.stdout.write_all(b"\n").await?; 34 | self.stdout.flush().await?; 35 | Ok(()) 36 | } 37 | } 38 | 39 | /// A mock writer implementation for testing. 40 | #[cfg(test)] 41 | pub struct MockWriter { 42 | /// Stores all messages written to this writer 43 | pub messages: Arc>>, 44 | } 45 | 46 | #[cfg(test)] 47 | impl MockWriter { 48 | pub fn new() -> Self { 49 | Self { 50 | messages: Arc::new(TokioMutex::new(Vec::new())), 51 | } 52 | } 53 | 54 | /// Get a clone of all messages that have been written 55 | pub async fn get_messages(&self) -> Vec { 56 | self.messages.lock().await.clone() 57 | } 58 | } 59 | 60 | #[cfg(test)] 61 | #[async_trait] 62 | impl MessageWriter for MockWriter { 63 | async fn write_message(&mut self, message: &str) -> Result<()> { 64 | let mut messages = self.messages.lock().await; 65 | messages.push(message.to_string()); 66 | Ok(()) 67 | } 68 | } 69 | 70 | #[cfg(test)] 71 | mod tests { 72 | use super::*; 73 | 74 | #[tokio::test] 75 | async fn test_stdout_writer() { 76 | // This is a simple test that just verifies the writer doesn't error 77 | // We can't easily test the actual output to stdout in a unit test 78 | let mut writer = StdoutWriter::new(tokio::io::stdout()); 79 | let result = writer.write_message("Test message").await; 80 | assert!(result.is_ok()); 81 | } 82 | 83 | #[tokio::test] 84 | async fn test_mock_writer() { 85 | let mut writer = MockWriter::new(); 86 | 87 | // Write some messages 88 | writer.write_message("Message 1").await.unwrap(); 89 | writer.write_message("Message 2").await.unwrap(); 90 | writer.write_message("Message 3").await.unwrap(); 91 | 92 | // Verify the messages were stored 93 | let messages = writer.get_messages().await; 94 | assert_eq!(messages.len(), 3); 95 | assert_eq!(messages[0], "Message 1"); 96 | assert_eq!(messages[1], "Message 2"); 97 | assert_eq!(messages[2], "Message 3"); 98 | 99 | // Clone the messages Arc for testing in another scope 100 | let messages_arc = writer.messages.clone(); 101 | 102 | // Ensure the message can be accessed from multiple places 103 | { 104 | let mut messages = messages_arc.lock().await; 105 | messages.push("Message 4".to_string()); 106 | } 107 | 108 | let messages = writer.get_messages().await; 109 | assert_eq!(messages.len(), 4); 110 | assert_eq!(messages[3], "Message 4"); 111 | } 112 | } 113 | -------------------------------------------------------------------------------- /crates/llm/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "llm" 3 | version = "0.1.6" 4 | edition = "2021" 5 | 6 | [dependencies] 7 | anyhow = "1.0" 8 | async-trait = "0.1" 9 | serde = { version = "1.0", features = ["derive"] } 10 | serde_json = "1.0" 11 | reqwest = { version = "0.11", features = ["json", "stream"] } 12 | tokio = { version = "1.44", features = ["full"] } 13 | futures = "0.3" 14 | tracing = "0.1" 15 | oauth2 = "4.4" 16 | base64 = "0.21" 17 | keyring = "2.3" 18 | chrono = { version = "0.4", features = ["serde"] } 19 | tempfile = "3.18" 20 | thiserror = "1.0" 21 | regex = "1.11" 22 | 23 | [dev-dependencies] 24 | axum = "0.7" 25 | bytes = "1.10" 26 | -------------------------------------------------------------------------------- /crates/llm/src/auth.rs: -------------------------------------------------------------------------------- 1 | use crate::config::DeploymentConfig; 2 | use anyhow::Result; 3 | use base64::engine::{general_purpose, Engine}; 4 | use std::sync::Arc; 5 | use std::time::{Duration, SystemTime}; 6 | use tokio::sync::RwLock; 7 | 8 | pub struct TokenManager { 9 | client_id: String, 10 | client_secret: String, 11 | token_url: String, 12 | current_token: RwLock>, 13 | } 14 | 15 | #[derive(serde::Deserialize)] 16 | struct TokenResponse { 17 | access_token: String, 18 | expires_in: u64, 19 | } 20 | 21 | struct TokenInfo { 22 | token: String, 23 | expires_at: SystemTime, 24 | } 25 | 26 | impl TokenManager { 27 | pub async fn new(config: &DeploymentConfig) -> Result> { 28 | tracing::debug!("Creating new TokenManager..."); 29 | 30 | let manager = Arc::new(Self { 31 | client_id: config.client_id.clone(), 32 | client_secret: config.client_secret.clone(), 33 | token_url: config.token_url.clone(), 34 | current_token: RwLock::new(None), 35 | }); 36 | 37 | // Fetch initial token 38 | manager.refresh_token().await?; 39 | 40 | Ok(manager) 41 | } 42 | 43 | pub async fn get_valid_token(&self) -> Result { 44 | // Check if we have a valid token 45 | if let Some(token_info) = self.current_token.read().await.as_ref() { 46 | if SystemTime::now() < token_info.expires_at { 47 | return Ok(token_info.token.clone()); 48 | } 49 | } 50 | 51 | // If not, we need to fetch a new one 52 | self.refresh_token().await 53 | } 54 | 55 | async fn refresh_token(&self) -> Result { 56 | tracing::debug!("Requesting new access token..."); 57 | 58 | let client = reqwest::Client::new(); 59 | let auth = 60 | general_purpose::STANDARD.encode(format!("{}:{}", self.client_id, self.client_secret)); 61 | 62 | let res = client 63 | .post(&self.token_url) 64 | .header("Authorization", format!("Basic {}", auth)) 65 | .header("Content-Type", "application/x-www-form-urlencoded") 66 | .body("grant_type=client_credentials") 67 | .send() 68 | .await?; 69 | 70 | let status = res.status(); 71 | if !status.is_success() { 72 | let error_text = res.text().await?; 73 | anyhow::bail!("Token request failed: {} - {}", status, error_text); 74 | } 75 | 76 | let token_response = res.json::().await?; 77 | 78 | // Set expiry slightly before actual expiry to ensure we don't use expired tokens 79 | let expires_at = SystemTime::now() + Duration::from_secs(token_response.expires_in - 60); 80 | 81 | let token_info = TokenInfo { 82 | token: token_response.access_token.clone(), 83 | expires_at, 84 | }; 85 | 86 | *self.current_token.write().await = Some(token_info); 87 | 88 | Ok(token_response.access_token) 89 | } 90 | } 91 | -------------------------------------------------------------------------------- /crates/llm/src/config.rs: -------------------------------------------------------------------------------- 1 | use anyhow::{Context, Result}; 2 | use keyring::Entry; 3 | use serde::{Deserialize, Serialize}; 4 | 5 | #[derive(Debug, Serialize, Deserialize)] 6 | pub struct DeploymentConfig { 7 | pub client_id: String, 8 | pub client_secret: String, 9 | pub token_url: String, 10 | pub api_base_url: String, 11 | } 12 | 13 | impl DeploymentConfig { 14 | const SERVICE_NAME: &'static str = "code-assistant-invoke"; 15 | const USERNAME: &'static str = "default"; 16 | 17 | pub fn load() -> Result { 18 | let keyring = Entry::new(Self::SERVICE_NAME, Self::USERNAME)?; 19 | 20 | match keyring.get_password() { 21 | Ok(config_json) => { 22 | serde_json::from_str(&config_json).with_context(|| "Failed to parse config") 23 | } 24 | Err(keyring::Error::NoEntry) => { 25 | // If no config exists, create it interactively 26 | let config = Self::create_interactive()?; 27 | config.save()?; 28 | Ok(config) 29 | } 30 | Err(e) => Err(e).with_context(|| "Failed to access keyring"), 31 | } 32 | } 33 | 34 | pub fn save(&self) -> Result<()> { 35 | let keyring = Entry::new(Self::SERVICE_NAME, Self::USERNAME)?; 36 | let config_json = 37 | serde_json::to_string(self).with_context(|| "Failed to serialize config")?; 38 | keyring 39 | .set_password(&config_json) 40 | .with_context(|| "Failed to save config to keyring") 41 | } 42 | 43 | fn create_interactive() -> Result { 44 | use std::io::{self, Write}; 45 | 46 | println!("No configuration found. Please enter the following details (they will be stored securely in your keyring):"); 47 | 48 | let mut input = String::new(); 49 | let mut config = DeploymentConfig { 50 | client_id: String::new(), 51 | client_secret: String::new(), 52 | token_url: String::new(), 53 | api_base_url: String::new(), 54 | }; 55 | 56 | print!("Client ID: "); 57 | io::stdout().flush().unwrap(); 58 | io::stdin().read_line(&mut input).unwrap(); 59 | // Remove whitespace but preserve special characters 60 | config.client_id = input.trim_end_matches(['\n', '\r']).to_string(); 61 | input.clear(); 62 | 63 | print!("Client Secret: "); 64 | io::stdout().flush().unwrap(); 65 | io::stdin().read_line(&mut input).unwrap(); 66 | // Same for secret 67 | config.client_secret = input.trim_end_matches(['\n', '\r']).to_string(); 68 | input.clear(); 69 | 70 | print!("Token URL: "); 71 | io::stdout().flush().unwrap(); 72 | io::stdin().read_line(&mut input).unwrap(); 73 | config.token_url = input.trim().to_string(); 74 | input.clear(); 75 | 76 | print!("API Base URL: "); 77 | io::stdout().flush().unwrap(); 78 | io::stdin().read_line(&mut input).unwrap(); 79 | config.api_base_url = input.trim().to_string(); 80 | input.clear(); 81 | 82 | Ok(config) 83 | } 84 | } 85 | -------------------------------------------------------------------------------- /crates/llm/src/display.rs: -------------------------------------------------------------------------------- 1 | use crate::types::{ContentBlock, Message, MessageContent}; 2 | use std::fmt; 3 | 4 | impl fmt::Display for ContentBlock { 5 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 6 | match self { 7 | ContentBlock::Text { text } => { 8 | writeln!(f, "Text: {}", text.replace('\n', "\n ")) 9 | } 10 | ContentBlock::ToolUse { id, name, input } => { 11 | writeln!(f, "ToolUse: id={}, name={}", id, name)?; 12 | writeln!( 13 | f, 14 | " Input: {}", 15 | serde_json::to_string_pretty(input) 16 | .unwrap_or_else(|_| input.to_string()) 17 | .replace('\n', "\n ") 18 | ) 19 | } 20 | ContentBlock::ToolResult { 21 | tool_use_id, 22 | content, 23 | is_error, 24 | } => { 25 | let error_suffix = if let Some(is_err) = is_error { 26 | if *is_err { 27 | " (ERROR)" 28 | } else { 29 | "" 30 | } 31 | } else { 32 | "" 33 | }; 34 | writeln!(f, "ToolResult: tool_use_id={}{}", tool_use_id, error_suffix)?; 35 | writeln!(f, " Content: {}", content.replace('\n', "\n ")) 36 | } 37 | ContentBlock::Thinking { 38 | thinking, 39 | signature, 40 | } => { 41 | writeln!(f, "Thinking: signature={}", signature)?; 42 | writeln!(f, " Content: {}", thinking.replace('\n', "\n ")) 43 | } 44 | ContentBlock::RedactedThinking { data } => { 45 | writeln!(f, "RedactedThinking")?; 46 | writeln!(f, " Data: {}", data.replace('\n', "\n ")) 47 | } 48 | } 49 | } 50 | } 51 | 52 | impl fmt::Display for MessageContent { 53 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 54 | match self { 55 | MessageContent::Text(content) => { 56 | writeln!(f, "Text: {}", content.replace('\n', "\n ")) 57 | } 58 | MessageContent::Structured(blocks) => { 59 | writeln!(f, "Structured content with {} blocks:", blocks.len())?; 60 | for (k, block) in blocks.iter().enumerate() { 61 | write!(f, " Block {}: ", k)?; 62 | // Convert the block display output to a string so we can add indentation 63 | let block_output = format!("{}", block); 64 | // Already includes a newline, so we don't need to add one here 65 | write!(f, "{}", block_output.replace('\n', "\n "))?; 66 | } 67 | Ok(()) 68 | } 69 | } 70 | } 71 | } 72 | 73 | impl fmt::Display for Message { 74 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 75 | writeln!(f, "Role: {:?}", self.role)?; 76 | write!(f, "{}", self.content) 77 | } 78 | } 79 | -------------------------------------------------------------------------------- /crates/llm/src/lib.rs: -------------------------------------------------------------------------------- 1 | //! LLM integration module providing abstraction over different LLM providers 2 | //! 3 | //! This module implements: 4 | //! - Common interface for LLM interactions via the LLMProvider trait 5 | //! - Support for multiple providers (Anthropic, OpenAI, Ollama, Vertex) 6 | //! - Message streaming capabilities 7 | //! - Provider-specific implementations and optimizations 8 | //! - Shared types and utilities for LLM interactions 9 | //! - Recording capabilities for debugging and testing 10 | 11 | #[cfg(test)] 12 | mod tests; 13 | 14 | mod utils; 15 | 16 | //pub mod aicore_converse; 17 | pub mod aicore_invoke; 18 | pub mod anthropic; 19 | pub mod anthropic_playback; 20 | pub mod auth; 21 | pub mod config; 22 | pub mod display; 23 | pub mod ollama; 24 | pub mod openai; 25 | pub mod openrouter; 26 | pub mod recording; 27 | pub mod types; 28 | pub mod vertex; 29 | 30 | pub use aicore_invoke::AiCoreClient; 31 | pub use anthropic::AnthropicClient; 32 | pub use ollama::OllamaClient; 33 | pub use openai::OpenAIClient; 34 | pub use openrouter::OpenRouterClient; 35 | pub use types::*; 36 | pub use vertex::{FixedToolIDGenerator, VertexClient}; 37 | 38 | use anyhow::Result; 39 | use async_trait::async_trait; 40 | 41 | /// Structure to represent different types of streaming content from LLMs 42 | #[derive(Debug, Clone)] 43 | pub enum StreamingChunk { 44 | /// Regular text content 45 | Text(String), 46 | /// Content identified as "thinking" (supported by some models) 47 | Thinking(String), 48 | /// JSON input for tool calls with optional metadata 49 | InputJson { 50 | content: String, 51 | tool_name: Option, 52 | tool_id: Option, 53 | }, 54 | } 55 | 56 | pub type StreamingCallback = Box Result<()> + Send + Sync>; 57 | 58 | /// Trait for different LLM provider implementations 59 | #[async_trait] 60 | pub trait LLMProvider { 61 | /// Sends a request to the LLM service 62 | async fn send_message( 63 | &self, 64 | request: LLMRequest, 65 | streaming_callback: Option<&StreamingCallback>, 66 | ) -> Result; 67 | } 68 | -------------------------------------------------------------------------------- /crates/llm/src/openrouter.rs: -------------------------------------------------------------------------------- 1 | use super::openai::OpenAIClient; 2 | use crate::{types::*, LLMProvider, StreamingCallback}; 3 | use anyhow::Result; 4 | use async_trait::async_trait; 5 | 6 | pub struct OpenRouterClient { 7 | inner: OpenAIClient, 8 | } 9 | 10 | impl OpenRouterClient { 11 | pub fn default_base_url() -> String { 12 | "https://openrouter.ai/api/v1".to_string() 13 | } 14 | 15 | pub fn new(api_key: String, model: String, base_url: String) -> Self { 16 | Self { 17 | inner: OpenAIClient::new(api_key, model, base_url), 18 | } 19 | } 20 | } 21 | 22 | #[async_trait] 23 | impl LLMProvider for OpenRouterClient { 24 | async fn send_message( 25 | &self, 26 | request: LLMRequest, 27 | streaming_callback: Option<&StreamingCallback>, 28 | ) -> Result { 29 | // Delegate to inner OpenAI client since the APIs are compatible 30 | self.inner.send_message(request, streaming_callback).await 31 | } 32 | } 33 | -------------------------------------------------------------------------------- /crates/llm/src/recording.rs: -------------------------------------------------------------------------------- 1 | use anyhow::{Context, Result}; 2 | use serde::{Deserialize, Serialize}; 3 | use std::fs::OpenOptions; 4 | use std::io::{Seek, Write}; 5 | use std::path::Path; 6 | use std::sync::{Arc, Mutex}; 7 | use std::time::Instant; 8 | 9 | /// Recording session that contains the original request and all chunks 10 | #[derive(Debug, Serialize, Deserialize, Clone)] 11 | pub struct RecordingSession { 12 | /// The request that was sent (simplified for storage) 13 | pub request: serde_json::Value, 14 | /// Timestamp of when the recording was started 15 | pub timestamp: chrono::DateTime, 16 | /// Raw chunks as received from the API 17 | pub chunks: Vec, 18 | } 19 | 20 | /// Single recorded chunk with timing info 21 | #[derive(Debug, Serialize, Deserialize, Clone)] 22 | pub struct RecordedChunk { 23 | /// Raw content of the data part of the SSE 24 | pub data: String, 25 | /// Milliseconds since recording start 26 | pub timestamp_ms: u64, 27 | } 28 | 29 | /// Recorder for API responses 30 | pub struct APIRecorder { 31 | file_path: Arc>>, 32 | current_session: Arc>>, 33 | start_time: Arc>>, 34 | } 35 | 36 | impl APIRecorder { 37 | /// Create a new recorder that writes to the specified file 38 | pub fn new>(path: P) -> Self { 39 | Self { 40 | file_path: Arc::new(Mutex::new(Some( 41 | path.as_ref().to_string_lossy().to_string(), 42 | ))), 43 | current_session: Arc::new(Mutex::new(None)), 44 | start_time: Arc::new(Mutex::new(None)), 45 | } 46 | } 47 | 48 | /// Start a new recording session 49 | pub fn start_recording(&self, request: serde_json::Value) -> Result<()> { 50 | let mut session_guard = self.current_session.lock().unwrap(); 51 | let mut start_guard = self.start_time.lock().unwrap(); 52 | 53 | // Create new session 54 | *session_guard = Some(RecordingSession { 55 | request, 56 | timestamp: chrono::Utc::now(), 57 | chunks: Vec::new(), 58 | }); 59 | 60 | // Record start time 61 | *start_guard = Some(Instant::now()); 62 | 63 | Ok(()) 64 | } 65 | 66 | /// Record an incoming chunk 67 | pub fn record_chunk(&self, data: &str) -> Result<()> { 68 | let mut session_guard = self.current_session.lock().unwrap(); 69 | let start_guard = self.start_time.lock().unwrap(); 70 | 71 | if let (Some(session), Some(start_time)) = (session_guard.as_mut(), *start_guard) { 72 | let elapsed = start_time.elapsed(); 73 | let timestamp_ms = elapsed.as_secs() * 1000 + elapsed.subsec_millis() as u64; 74 | 75 | session.chunks.push(RecordedChunk { 76 | data: data.to_string(), 77 | timestamp_ms, 78 | }); 79 | } 80 | 81 | Ok(()) 82 | } 83 | 84 | /// End the current recording session and save to disk 85 | pub fn end_recording(&self) -> Result<()> { 86 | let file_path_guard = self.file_path.lock().unwrap(); 87 | let mut session_guard = self.current_session.lock().unwrap(); 88 | let mut start_guard = self.start_time.lock().unwrap(); 89 | 90 | if let (Some(file_path), Some(session)) = (file_path_guard.as_ref(), session_guard.take()) { 91 | // Create/open the file 92 | let mut file = OpenOptions::new() 93 | .create(true) 94 | .append(true) 95 | .open(file_path) 96 | .context("Failed to open recording file")?; 97 | 98 | // Serialize and write the session 99 | let json = serde_json::to_string_pretty(&session)?; 100 | if let Ok(metadata) = std::fs::metadata(file_path) { 101 | let file_size = metadata.len(); 102 | 103 | if file_size == 0 { 104 | // If file is empty, start a JSON array 105 | writeln!(file, "[")?; 106 | } else { 107 | // If file already has content, add a comma 108 | // Go to position after the last } bracket, 109 | // i.e. skip "\n]\n" backwards from the end of the file 110 | file.set_len(file_size - 3)?; 111 | file.seek(std::io::SeekFrom::End(0))?; 112 | writeln!(file, ",")?; 113 | } 114 | } 115 | 116 | // Write the session 117 | let mut file = OpenOptions::new() 118 | .append(true) 119 | .open(file_path) 120 | .context("Failed to open recording file")?; 121 | 122 | writeln!(file, "{}", json)?; 123 | writeln!(file, "]")?; 124 | } 125 | 126 | // Reset start time 127 | *start_guard = None; 128 | 129 | Ok(()) 130 | } 131 | } 132 | -------------------------------------------------------------------------------- /crates/llm/src/types.rs: -------------------------------------------------------------------------------- 1 | use reqwest::Response; 2 | use serde::{Deserialize, Serialize}; 3 | use std::time::Duration; 4 | 5 | /// Tracks token usage for a request/response pair 6 | #[derive(Debug, Deserialize, PartialEq, Clone, Default)] 7 | pub struct Usage { 8 | /// Number of tokens in the input (prompt) 9 | pub input_tokens: u32, 10 | /// Number of tokens in the output (completion) 11 | pub output_tokens: u32, 12 | /// Number of tokens written to cache 13 | #[serde(default)] 14 | pub cache_creation_input_tokens: u32, 15 | /// Number of tokens read from cache 16 | #[serde(default)] 17 | pub cache_read_input_tokens: u32, 18 | } 19 | 20 | impl Usage { 21 | pub fn zero() -> Self { 22 | Usage { 23 | input_tokens: 0, 24 | output_tokens: 0, 25 | cache_creation_input_tokens: 0, 26 | cache_read_input_tokens: 0, 27 | } 28 | } 29 | } 30 | 31 | #[derive(Debug, Clone, Serialize)] 32 | pub struct ToolDefinition { 33 | pub name: String, 34 | pub description: String, 35 | pub parameters: serde_json::Value, 36 | } 37 | 38 | /// Generic request structure that can be mapped to different providers 39 | #[derive(Debug, Clone, Default)] 40 | pub struct LLMRequest { 41 | pub messages: Vec, 42 | pub system_prompt: String, 43 | pub tools: Option>, 44 | } 45 | 46 | #[derive(Debug, Serialize, Deserialize, Clone)] 47 | pub struct Message { 48 | pub role: MessageRole, 49 | pub content: MessageContent, 50 | } 51 | 52 | #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] 53 | #[serde(rename_all = "lowercase")] 54 | pub enum MessageRole { 55 | User, 56 | Assistant, 57 | } 58 | 59 | #[derive(Debug, Serialize, Deserialize, Clone)] 60 | #[serde(untagged)] 61 | pub enum MessageContent { 62 | Text(String), 63 | Structured(Vec), 64 | } 65 | 66 | #[derive(Debug, Serialize, Deserialize, PartialEq, Clone)] 67 | #[serde(tag = "type")] 68 | pub enum ContentBlock { 69 | #[serde(rename = "thinking")] 70 | Thinking { thinking: String, signature: String }, 71 | 72 | #[serde(rename = "redacted_thinking")] 73 | RedactedThinking { data: String }, 74 | 75 | #[serde(rename = "text")] 76 | Text { text: String }, 77 | 78 | #[serde(rename = "tool_use")] 79 | ToolUse { 80 | id: String, 81 | name: String, 82 | input: serde_json::Value, 83 | }, 84 | 85 | #[serde(rename = "tool_result")] 86 | ToolResult { 87 | tool_use_id: String, 88 | content: String, 89 | #[serde(skip_serializing_if = "Option::is_none")] 90 | is_error: Option, 91 | }, 92 | } 93 | 94 | /// Generic response structure 95 | #[derive(Debug, Deserialize, Clone, Default)] 96 | pub struct LLMResponse { 97 | pub content: Vec, 98 | pub usage: Usage, 99 | } 100 | 101 | /// Common error types for all LLM providers 102 | #[derive(Debug, thiserror::Error)] 103 | pub enum ApiError { 104 | #[error("Rate limit exceeded: {0}")] 105 | RateLimit(String), 106 | 107 | #[error("Authentication failed: {0}")] 108 | Authentication(String), 109 | 110 | #[error("Invalid request: {0}")] 111 | InvalidRequest(String), 112 | 113 | #[error("Service error: {0}")] 114 | ServiceError(String), 115 | 116 | #[error("Service overloaded: {0}")] 117 | Overloaded(String), 118 | 119 | #[error("Network error: {0}")] 120 | NetworkError(String), 121 | 122 | #[error("Unknown error: {0}")] 123 | Unknown(String), 124 | } 125 | 126 | /// Context wrapper for API errors that includes rate limit information 127 | #[derive(Debug, thiserror::Error)] 128 | #[error("{error}")] 129 | pub struct ApiErrorContext { 130 | pub error: ApiError, 131 | pub rate_limits: Option, 132 | } 133 | 134 | /// Base trait for rate limit information 135 | pub trait RateLimitHandler: Sized { 136 | /// Create a new instance from response headers 137 | fn from_response(response: &Response) -> Self; 138 | 139 | /// Get the delay duration before the next retry 140 | fn get_retry_delay(&self) -> Duration; 141 | 142 | /// Log the current rate limit status 143 | fn log_status(&self); 144 | } 145 | -------------------------------------------------------------------------------- /crates/llm/src/utils.rs: -------------------------------------------------------------------------------- 1 | use crate::{ApiError, ApiErrorContext, RateLimitHandler}; 2 | use anyhow::Result; 3 | use reqwest::{Response, StatusCode}; 4 | use std::time::Duration; 5 | use tokio::time::sleep; 6 | use tracing::warn; 7 | 8 | /// Check response error and extract rate limit information. 9 | /// Returns Ok(Response) if successful, or an error with rate limit context if not. 10 | pub async fn check_response_error( 11 | response: Response, 12 | ) -> Result { 13 | let status = response.status(); 14 | if status.is_success() { 15 | return Ok(response); 16 | } 17 | 18 | let rate_limits = T::from_response(&response); 19 | let response_text = response 20 | .text() 21 | .await 22 | .map_err(|e| ApiError::NetworkError(e.to_string()))?; 23 | 24 | let error = match status { 25 | StatusCode::TOO_MANY_REQUESTS => ApiError::RateLimit(response_text), 26 | StatusCode::UNAUTHORIZED => ApiError::Authentication(response_text), 27 | StatusCode::BAD_REQUEST => ApiError::InvalidRequest(response_text), 28 | status if status.is_server_error() => ApiError::ServiceError(response_text), 29 | _ => ApiError::Unknown(format!("Status {}: {}", status, response_text)), 30 | }; 31 | 32 | Err(ApiErrorContext { 33 | error, 34 | rate_limits: Some(rate_limits), 35 | } 36 | .into()) 37 | } 38 | 39 | /// Handle retryable errors and rate limiting for LLM providers. 40 | /// Returns true if the error is retryable and we should continue the retry loop. 41 | /// Returns false if we should exit the retry loop. 42 | pub async fn handle_retryable_error< 43 | T: RateLimitHandler + std::fmt::Debug + Send + Sync + 'static, 44 | >( 45 | error: &anyhow::Error, 46 | attempts: u32, 47 | max_retries: u32, 48 | ) -> bool { 49 | if let Some(ctx) = error.downcast_ref::>() { 50 | match &ctx.error { 51 | ApiError::RateLimit(_) => { 52 | if let Some(rate_limits) = &ctx.rate_limits { 53 | if attempts < max_retries { 54 | let delay = rate_limits.get_retry_delay(); 55 | warn!( 56 | "Rate limit hit (attempt {}/{}), waiting {} seconds before retry", 57 | attempts, 58 | max_retries, 59 | delay.as_secs() 60 | ); 61 | sleep(delay).await; 62 | return true; 63 | } 64 | } else { 65 | // Fallback if no rate limit info available 66 | if attempts < max_retries { 67 | let delay = Duration::from_secs(2u64.pow(attempts - 1)); 68 | warn!( 69 | "Rate limit hit but no timing info available (attempt {}/{}), using exponential backoff: {} seconds", 70 | attempts, 71 | max_retries, 72 | delay.as_secs() 73 | ); 74 | sleep(delay).await; 75 | return true; 76 | } 77 | } 78 | } 79 | ApiError::ServiceError(_) | ApiError::NetworkError(_) | ApiError::Overloaded(_) => { 80 | if attempts < max_retries { 81 | let delay = Duration::from_secs(2u64.pow(attempts - 1)); 82 | warn!( 83 | "Error: {} (attempt {}/{}), retrying in {} seconds", 84 | error, 85 | attempts, 86 | max_retries, 87 | delay.as_secs() 88 | ); 89 | sleep(delay).await; 90 | return true; 91 | } 92 | } 93 | _ => { 94 | warn!( 95 | "Unhandled error (attempt {}/{}): {:?}", 96 | attempts, max_retries, error 97 | ); 98 | } 99 | } 100 | } 101 | false 102 | } 103 | -------------------------------------------------------------------------------- /crates/web/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "web" 3 | version = "0.1.6" 4 | edition = "2021" 5 | 6 | [dependencies] 7 | anyhow = "1.0" 8 | async-trait = "0.1" 9 | chromiumoxide = { version = "0.5", features = ["tokio-runtime"] } 10 | futures = "0.3" 11 | htmd = "0.1.6" 12 | percent-encoding = "2.3" 13 | regex = "1.11" 14 | reqwest = { version = "0.11", features = ["json", "stream"] } 15 | scraper = "0.18" 16 | serde = { version = "1.0", features = ["derive"] } 17 | serde_json = "1.0" 18 | tempfile = "3.18" 19 | tokio = { version = "1.44", features = ["full"] } 20 | url = "2.5" 21 | 22 | [dev-dependencies] 23 | axum = "0.7" 24 | -------------------------------------------------------------------------------- /crates/web/src/lib.rs: -------------------------------------------------------------------------------- 1 | mod client; 2 | mod perplexity; 3 | #[cfg(test)] 4 | mod tests; 5 | pub use client::{PageMetadata, WebClient, WebPage, WebSearchResult}; 6 | pub use perplexity::{PerplexityCitation, PerplexityClient, PerplexityMessage, PerplexityResponse}; 7 | -------------------------------------------------------------------------------- /crates/web/src/tests.rs: -------------------------------------------------------------------------------- 1 | #[cfg(test)] 2 | use super::WebClient; 3 | 4 | #[tokio::test] 5 | async fn test_web_search() { 6 | let client = WebClient::new().await.unwrap(); 7 | let results = client.search("rust programming", 1).await.unwrap(); 8 | 9 | println!("\nSearch Results:"); 10 | for (i, result) in results.iter().enumerate() { 11 | println!("\n{}. {}", i + 1, result.title); 12 | println!(" URL: {}", result.url); 13 | println!(" Snippet: {}", result.snippet); 14 | } 15 | 16 | assert!(!results.is_empty()); 17 | assert!(!results[0].url.is_empty()); 18 | assert!(!results[0].title.is_empty()); 19 | assert!(!results[0].snippet.is_empty()); 20 | } 21 | 22 | #[tokio::test] 23 | async fn test_web_fetch() { 24 | let client = WebClient::new().await.unwrap(); 25 | let page = client.fetch("https://www.rust-lang.org").await.unwrap(); 26 | 27 | println!("\nContent: {}", page.content); 28 | 29 | assert!(!page.content.is_empty()); 30 | assert!(page.content.contains("Rust")); 31 | } 32 | --------------------------------------------------------------------------------