Compare commits
24
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
f44722df04
|
||
|
|
3eaae0e652 | ||
|
|
b12829db39 | ||
|
|
dffaf6b9db | ||
|
|
a9a4ccca88 | ||
|
|
7cb7d66575 | ||
|
|
400b50fbd0 | ||
|
|
7eeff2a226 | ||
|
|
374e8c0bf7
|
||
|
|
ede87df960
|
||
|
|
95a8c3df44
|
||
|
|
f57bd21ee4 | ||
|
|
e1b6e3f8c6 | ||
|
|
644d899f78 | ||
|
|
2596194417 | ||
|
|
684f19250a | ||
|
|
f3d59ade11 | ||
|
|
e1604c58ea | ||
|
|
dcacb3a962 | ||
|
|
d791098e51 | ||
|
|
d31110cd67
|
||
|
|
0f35e03a85
|
||
|
|
c2b0c120d7
|
||
|
|
65c9be36b2
|
+4
-1
@@ -5,4 +5,7 @@
|
|||||||
.idea/
|
.idea/
|
||||||
/coyote.iml
|
/coyote.iml
|
||||||
/.idea/
|
/.idea/
|
||||||
.coyote
|
.coyote/**
|
||||||
|
.sisyphus/**
|
||||||
|
.coyote-project.json
|
||||||
|
.coyote/memory/
|
||||||
|
|||||||
Generated
+43
-8
@@ -1962,6 +1962,16 @@ dependencies = [
|
|||||||
"darling_macro 0.23.0",
|
"darling_macro 0.23.0",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "darling"
|
||||||
|
version = "0.24.0"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "88490bf1b990d87eaaa7ac8aa887f629a08e7359765b4911faf63c3763347d23"
|
||||||
|
dependencies = [
|
||||||
|
"darling_core 0.24.0",
|
||||||
|
"darling_macro 0.24.0",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "darling_core"
|
name = "darling_core"
|
||||||
version = "0.20.11"
|
version = "0.20.11"
|
||||||
@@ -1989,6 +1999,19 @@ dependencies = [
|
|||||||
"syn 2.0.119",
|
"syn 2.0.119",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "darling_core"
|
||||||
|
version = "0.24.0"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "084e274f91c482280130e1e34e0b8d6e66776a060d7b6de7b84289ca778868c4"
|
||||||
|
dependencies = [
|
||||||
|
"ident_case",
|
||||||
|
"proc-macro2",
|
||||||
|
"quote",
|
||||||
|
"strsim",
|
||||||
|
"syn 3.0.3",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "darling_macro"
|
name = "darling_macro"
|
||||||
version = "0.20.11"
|
version = "0.20.11"
|
||||||
@@ -2011,6 +2034,17 @@ dependencies = [
|
|||||||
"syn 2.0.119",
|
"syn 2.0.119",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "darling_macro"
|
||||||
|
version = "0.24.0"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "68f5792fa0d41cd2325ce0ffa64f0a340eaebd4971a3a0c5e1ffd2cc488a355e"
|
||||||
|
dependencies = [
|
||||||
|
"darling_core 0.24.0",
|
||||||
|
"quote",
|
||||||
|
"syn 3.0.3",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "defmt"
|
name = "defmt"
|
||||||
version = "1.1.1"
|
version = "1.1.1"
|
||||||
@@ -5303,12 +5337,12 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "rmcp"
|
name = "rmcp"
|
||||||
version = "1.8.0"
|
version = "3.1.2"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "1d1f571c72940a19d9532fe52dbea8bc9912bf1d766c2970bb824056b86f3f59"
|
checksum = "c8dddc5b1924b9a59fba420166160ca2c4663a4e01803e52eda33070f56d63c8"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"async-trait",
|
"base64 0.23.1",
|
||||||
"base64 0.22.1",
|
"bytes",
|
||||||
"chrono",
|
"chrono",
|
||||||
"futures",
|
"futures",
|
||||||
"http 1.5.0",
|
"http 1.5.0",
|
||||||
@@ -5326,19 +5360,20 @@ dependencies = [
|
|||||||
"tokio-stream",
|
"tokio-stream",
|
||||||
"tokio-util",
|
"tokio-util",
|
||||||
"tracing",
|
"tracing",
|
||||||
|
"uuid",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "rmcp-macros"
|
name = "rmcp-macros"
|
||||||
version = "1.8.0"
|
version = "3.1.2"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "1aad0035b69380782d78ea95b508327e6deaa2235909053e596eea8f27b5e1d5"
|
checksum = "6898e24cd16342b59bfa8a53c2c04b9cf62fc8a2cfea57b9c038b09984bfc521"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"darling 0.23.0",
|
"darling 0.24.0",
|
||||||
"proc-macro2",
|
"proc-macro2",
|
||||||
"quote",
|
"quote",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
"syn 2.0.119",
|
"syn 3.0.3",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
|
|||||||
+2
-2
@@ -84,7 +84,7 @@ duct = "1.0.0"
|
|||||||
argc = "1.23.0"
|
argc = "1.23.0"
|
||||||
strum_macros = "0.27.2"
|
strum_macros = "0.27.2"
|
||||||
indoc = "2.0.6"
|
indoc = "2.0.6"
|
||||||
rmcp = { version = "1.5.0", features = [
|
rmcp = { version = "3.1.2", features = [
|
||||||
"client",
|
"client",
|
||||||
"transport-child-process",
|
"transport-child-process",
|
||||||
"transport-streamable-http-client-reqwest",
|
"transport-streamable-http-client-reqwest",
|
||||||
@@ -150,4 +150,4 @@ path = "src/main.rs"
|
|||||||
[profile.release]
|
[profile.release]
|
||||||
lto = true
|
lto = true
|
||||||
strip = true
|
strip = true
|
||||||
opt-level = "z"
|
opt-level = "z"
|
||||||
@@ -31,7 +31,7 @@ settings:
|
|||||||
max_loop_iterations: 20
|
max_loop_iterations: 20
|
||||||
log_state_snapshots: true
|
log_state_snapshots: true
|
||||||
validate_before_run: true
|
validate_before_run: true
|
||||||
timeout: 1800
|
timeout: 14400
|
||||||
|
|
||||||
initial_state:
|
initial_state:
|
||||||
project_dir: ''
|
project_dir: ''
|
||||||
@@ -90,6 +90,7 @@ nodes:
|
|||||||
Project directory: {{project_dir}}
|
Project directory: {{project_dir}}
|
||||||
prompt: '{{initial_prompt}}'
|
prompt: '{{initial_prompt}}'
|
||||||
tools: []
|
tools: []
|
||||||
|
timeout: 300
|
||||||
output_schema:
|
output_schema:
|
||||||
type: object
|
type: object
|
||||||
properties:
|
properties:
|
||||||
@@ -254,6 +255,7 @@ nodes:
|
|||||||
- fs_patch
|
- fs_patch
|
||||||
- execute_command
|
- execute_command
|
||||||
max_iterations: 100
|
max_iterations: 100
|
||||||
|
timeout: 1800
|
||||||
state_updates:
|
state_updates:
|
||||||
last_node_output: '{{output}}'
|
last_node_output: '{{output}}'
|
||||||
fallback: end_failure
|
fallback: end_failure
|
||||||
@@ -327,6 +329,7 @@ nodes:
|
|||||||
- fs_ls
|
- fs_ls
|
||||||
- execute_command
|
- execute_command
|
||||||
max_iterations: 15
|
max_iterations: 15
|
||||||
|
timeout: 600
|
||||||
output_schema:
|
output_schema:
|
||||||
type: object
|
type: object
|
||||||
properties:
|
properties:
|
||||||
@@ -384,4 +387,4 @@ nodes:
|
|||||||
{{build_output}}
|
{{build_output}}
|
||||||
|
|
||||||
Last tests output:
|
Last tests output:
|
||||||
{{tests_output}}
|
{{tests_output}}
|
||||||
@@ -14,6 +14,9 @@ variables:
|
|||||||
- name: project_dir
|
- name: project_dir
|
||||||
description: Absolute path to the project the plan targets - the ground truth for pointer verification
|
description: Absolute path to the project the plan targets - the ground truth for pointer verification
|
||||||
default: '.'
|
default: '.'
|
||||||
|
- name: auto_confirm
|
||||||
|
description: Auto-confirm command execution
|
||||||
|
default: '1'
|
||||||
|
|
||||||
global_tools:
|
global_tools:
|
||||||
- ast_grep.sh
|
- ast_grep.sh
|
||||||
|
|||||||
@@ -21,8 +21,7 @@
|
|||||||
},
|
},
|
||||||
"iwe": {
|
"iwe": {
|
||||||
"type": "stdio",
|
"type": "stdio",
|
||||||
"command": "iwec",
|
"command": "iwec"
|
||||||
"args": ["--project", "."]
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -14,6 +14,11 @@ set -e
|
|||||||
# is the most common cause of "unable to apply patch" failures, especially in files with sed/jq/regex pipelines or
|
# is the most common cause of "unable to apply patch" failures, especially in files with sed/jq/regex pipelines or
|
||||||
# embedded Python with quoted strings.
|
# embedded Python with quoted strings.
|
||||||
# - Hunks are applied in order; the first hunk that fails aborts the whole patch — later hunks are NOT attempted.
|
# - Hunks are applied in order; the first hunk that fails aborts the whole patch — later hunks are NOT attempted.
|
||||||
|
# - Hunks anchor at the FIRST exact match of their context in the file, so include enough context to make each hunk
|
||||||
|
# unique. Hunks must appear in file order.
|
||||||
|
# - Every hunk needs at least one context or removed line to anchor it; a hunk containing only additions is an error.
|
||||||
|
# - Unrecognized lines inside a hunk are hard errors: every hunk line must start with ' ' (context), '-' (removal),
|
||||||
|
# or '+' (addition).
|
||||||
# - If you've edited this file in earlier tool calls, fs_cat it again before composing the patch. A stale view of the file
|
# - If you've edited this file in earlier tool calls, fs_cat it again before composing the patch. A stale view of the file
|
||||||
# produces context lines that no longer match.
|
# produces context lines that no longer match.
|
||||||
# - On failure the error message names the failing hunk and shows the expected-vs-actual line. Fix that specific line and
|
# - On failure the error message names the failing hunk and shows the expected-vs-actual line. Fix that specific line and
|
||||||
@@ -33,7 +38,11 @@ source "$LLM_PROMPT_UTILS_FILE"
|
|||||||
|
|
||||||
# shellcheck disable=SC2154
|
# shellcheck disable=SC2154
|
||||||
main() {
|
main() {
|
||||||
argc_contents="$(jq -r '.content' <<< "$LLM_TOOL_RAW_JSON")"
|
# Command substitution strips *all* trailing newlines and `jq -r` appends one
|
||||||
|
# of its own, so read with `-j` and pin the real end of the content with a
|
||||||
|
# sentinel that is removed afterwards.
|
||||||
|
argc_contents="$(jq -j '.content' <<< "$LLM_TOOL_RAW_JSON"; printf x)"
|
||||||
|
argc_contents="${argc_contents%x}"
|
||||||
argc_path="$(jq -r '.path' <<< "$LLM_TOOL_RAW_JSON")"
|
argc_path="$(jq -r '.path' <<< "$LLM_TOOL_RAW_JSON")"
|
||||||
|
|
||||||
if [[ ! -f "$argc_path" ]]; then
|
if [[ ! -f "$argc_path" ]]; then
|
||||||
@@ -41,7 +50,19 @@ main() {
|
|||||||
exit 1
|
exit 1
|
||||||
fi
|
fi
|
||||||
|
|
||||||
new_contents="$(patch_file "$argc_path" <(printf "%s" "$argc_contents"))"
|
# Same sentinel guard on the patched result, otherwise the trailing newline
|
||||||
|
# is stripped again on the way back out. `rc` preserves patch_file's exit
|
||||||
|
# status so a failure still aborts under `set -e`.
|
||||||
|
new_contents="$(patch_file "$argc_path" <(printf "%s" "$argc_contents"); rc=$?; printf x; exit "$rc")"
|
||||||
|
new_contents="${new_contents%x}"
|
||||||
|
|
||||||
|
# awk newline-terminates every printed line, so patching a file that lacks
|
||||||
|
# a final newline would silently add one. Preserve the original file's
|
||||||
|
# final-newline state instead.
|
||||||
|
if [[ -n "$(tail -c 1 "$argc_path")" ]]; then
|
||||||
|
new_contents="${new_contents%$'\n'}"
|
||||||
|
fi
|
||||||
|
|
||||||
printf "%s" "$new_contents" | git diff --no-index "$argc_path" - || true
|
printf "%s" "$new_contents" | git diff --no-index "$argc_path" - || true
|
||||||
|
|
||||||
guard_operation "Apply changes?"
|
guard_operation "Apply changes?"
|
||||||
|
|||||||
@@ -15,7 +15,12 @@ source "$LLM_PROMPT_UTILS_FILE"
|
|||||||
|
|
||||||
# shellcheck disable=SC2154
|
# shellcheck disable=SC2154
|
||||||
main() {
|
main() {
|
||||||
argc_contents="$(jq -r '.content' <<< "$LLM_TOOL_RAW_JSON")"
|
# Command substitution strips *all* trailing newlines and `jq -r` appends one
|
||||||
|
# of its own, so read with `-j` and pin the real end of the content with a
|
||||||
|
# sentinel that is removed afterwards. Without this every written file loses
|
||||||
|
# its final newline, which breaks formatters such as `cargo fmt --check`.
|
||||||
|
argc_contents="$(jq -j '.content' <<< "$LLM_TOOL_RAW_JSON"; printf x)"
|
||||||
|
argc_contents="${argc_contents%x}"
|
||||||
argc_path="$(jq -r '.path' <<< "$LLM_TOOL_RAW_JSON")"
|
argc_path="$(jq -r '.path' <<< "$LLM_TOOL_RAW_JSON")"
|
||||||
|
|
||||||
if [[ -f "$argc_path" ]]; then
|
if [[ -f "$argc_path" ]]; then
|
||||||
|
|||||||
@@ -186,7 +186,9 @@ input() {
|
|||||||
}
|
}
|
||||||
|
|
||||||
confirm() {
|
confirm() {
|
||||||
trap "stty echo; exit" EXIT
|
# stty targets stdin, which is /dev/null when the host spawns tool scripts;
|
||||||
|
# point it at the real terminal and stay quiet when there isn't one.
|
||||||
|
trap "stty echo </dev/tty 2>/dev/null; exit" EXIT
|
||||||
_prompt_text "$1 (y/N)"
|
_prompt_text "$1 (y/N)"
|
||||||
echo -en "\033[36m\c " >&2
|
echo -en "\033[36m\c " >&2
|
||||||
|
|
||||||
@@ -229,7 +231,7 @@ list() {
|
|||||||
declare first_row
|
declare first_row
|
||||||
first_row=$((last_row - opts_count + 1))
|
first_row=$((last_row - opts_count + 1))
|
||||||
|
|
||||||
trap "_cursor_blink_on; stty echo; exit" 2
|
trap "_cursor_blink_on; stty echo </dev/tty 2>/dev/null; exit" 2
|
||||||
|
|
||||||
_cursor_blink_off
|
_cursor_blink_off
|
||||||
|
|
||||||
@@ -275,7 +277,7 @@ checkbox() {
|
|||||||
declare first_row
|
declare first_row
|
||||||
first_row=$((last_row - opts_count + 1))
|
first_row=$((last_row - opts_count + 1))
|
||||||
|
|
||||||
trap "_cursor_blink_on; stty echo; exit" 2
|
trap "_cursor_blink_on; stty echo </dev/tty 2>/dev/null; exit" 2
|
||||||
|
|
||||||
_cursor_blink_off
|
_cursor_blink_off
|
||||||
|
|
||||||
@@ -403,7 +405,7 @@ range() {
|
|||||||
declare current_row
|
declare current_row
|
||||||
current_row=$((first_row - 1))
|
current_row=$((first_row - 1))
|
||||||
|
|
||||||
trap "_cursor_blink_on; stty echo; exit" 2
|
trap "_cursor_blink_on; stty echo </dev/tty 2>/dev/null; exit" 2
|
||||||
|
|
||||||
_cursor_blink_off
|
_cursor_blink_off
|
||||||
|
|
||||||
@@ -528,7 +530,11 @@ guard_operation() {
|
|||||||
# + print(f"Hello {name}")
|
# + print(f"Hello {name}")
|
||||||
patch_file() {
|
patch_file() {
|
||||||
awk '
|
awk '
|
||||||
FNR == NR {
|
function isHeaderPair(i) {
|
||||||
|
return (patchLines[i] ~ /^--- / && patchLines[i+1] ~ /^\+\+\+ / && patchLines[i+2] ~ /^@@/)
|
||||||
|
}
|
||||||
|
|
||||||
|
FILENAME == ARGV[1] {
|
||||||
lines[FNR] = $0
|
lines[FNR] = $0
|
||||||
next;
|
next;
|
||||||
}
|
}
|
||||||
@@ -547,11 +553,6 @@ patch_file() {
|
|||||||
while (patchLineIndex <= totalPatchLines) {
|
while (patchLineIndex <= totalPatchLines) {
|
||||||
line = patchLines[patchLineIndex]
|
line = patchLines[patchLineIndex]
|
||||||
|
|
||||||
if (line ~ /^--- / || line ~ /^\+\+\+ /) {
|
|
||||||
patchLineIndex++
|
|
||||||
continue
|
|
||||||
}
|
|
||||||
|
|
||||||
if (line ~ /^@@/) {
|
if (line ~ /^@@/) {
|
||||||
mode = "hunk"
|
mode = "hunk"
|
||||||
hunkIndex++
|
hunkIndex++
|
||||||
@@ -560,7 +561,22 @@ patch_file() {
|
|||||||
}
|
}
|
||||||
|
|
||||||
if (mode == "hunk") {
|
if (mode == "hunk") {
|
||||||
while (patchLineIndex <= totalPatchLines && line ~ /^[-+ ]|^\s*$/ && line !~ /^--- /) {
|
while (patchLineIndex <= totalPatchLines) {
|
||||||
|
line = patchLines[patchLineIndex]
|
||||||
|
|
||||||
|
if (line ~ /^\\ No newline/) {
|
||||||
|
patchLineIndex++
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
if (isHeaderPair(patchLineIndex)) {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
|
||||||
|
if (line !~ /^[-+ ]/ && line !~ /^[ \t]*$/) {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
|
||||||
sanitizedLine = substr(line, 2)
|
sanitizedLine = substr(line, 2)
|
||||||
|
|
||||||
if (line !~ /^\+/) {
|
if (line !~ /^\+/) {
|
||||||
@@ -574,13 +590,37 @@ patch_file() {
|
|||||||
}
|
}
|
||||||
|
|
||||||
patchLineIndex++
|
patchLineIndex++
|
||||||
line = patchLines[patchLineIndex]
|
|
||||||
}
|
}
|
||||||
|
|
||||||
mode = "none"
|
mode = "none"
|
||||||
} else {
|
continue
|
||||||
patchLineIndex++
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if (isHeaderPair(patchLineIndex)) {
|
||||||
|
patchLineIndex += 2
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
if (line ~ /^\\ No newline/) {
|
||||||
|
patchLineIndex++
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
if (hunkIndex == 0) {
|
||||||
|
# Preamble before the first hunk: tolerate prose, code fences, and lone headers.
|
||||||
|
patchLineIndex++
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
if (line ~ /^[ \t]*$/ || line ~ /^```/) {
|
||||||
|
patchLineIndex++
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
print "error: unrecognized line in patch (line " patchLineIndex "): " line > "/dev/stderr"
|
||||||
|
print "" > "/dev/stderr"
|
||||||
|
print "Every line inside a hunk must start with \" \" (context), \"-\" (removal), or \"+\" (addition)." > "/dev/stderr"
|
||||||
|
exit 1
|
||||||
}
|
}
|
||||||
|
|
||||||
if (hunkIndex == 0) {
|
if (hunkIndex == 0) {
|
||||||
@@ -593,6 +633,36 @@ patch_file() {
|
|||||||
}
|
}
|
||||||
|
|
||||||
totalHunks = hunkIndex
|
totalHunks = hunkIndex
|
||||||
|
|
||||||
|
if (totalLines == 0) {
|
||||||
|
for (h = 1; h <= totalHunks; h++) {
|
||||||
|
if (hunkTotalOriginalLines[h] > 0) {
|
||||||
|
print "error: unable to apply patch" > "/dev/stderr"
|
||||||
|
print "" > "/dev/stderr"
|
||||||
|
print "Hunk " h " expects existing content but the file is empty." > "/dev/stderr"
|
||||||
|
exit 1
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
for (h = 1; h <= totalHunks; h++) {
|
||||||
|
for (i = 1; i <= hunkTotalUpdatedLines[h]; i++) {
|
||||||
|
print hunkUpdatedLines[h,i]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
exit 0
|
||||||
|
}
|
||||||
|
|
||||||
|
for (h = 1; h <= totalHunks; h++) {
|
||||||
|
if (hunkTotalOriginalLines[h] == 0) {
|
||||||
|
print "error: unable to apply patch" > "/dev/stderr"
|
||||||
|
print "" > "/dev/stderr"
|
||||||
|
print "Hunk " h " contains no context or removed lines; include at least one" > "/dev/stderr"
|
||||||
|
print "context line so the hunk can be anchored." > "/dev/stderr"
|
||||||
|
exit 1
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
hunkIndex = 1
|
hunkIndex = 1
|
||||||
|
|
||||||
for (lineIndex = 1; lineIndex <= totalLines; lineIndex++) {
|
for (lineIndex = 1; lineIndex <= totalLines; lineIndex++) {
|
||||||
@@ -603,7 +673,7 @@ patch_file() {
|
|||||||
nextLineIndex = lineIndex + 1
|
nextLineIndex = lineIndex + 1
|
||||||
|
|
||||||
for (i = 2; i <= hunkTotalOriginalLines[hunkIndex]; i++) {
|
for (i = 2; i <= hunkTotalOriginalLines[hunkIndex]; i++) {
|
||||||
if (lines[nextLineIndex] != hunkOriginalLines[hunkIndex,i]) {
|
if (nextLineIndex > totalLines || lines[nextLineIndex] != hunkOriginalLines[hunkIndex,i]) {
|
||||||
if (i - 1 > bestPartialLen[hunkIndex]) {
|
if (i - 1 > bestPartialLen[hunkIndex]) {
|
||||||
bestPartialLen[hunkIndex] = i - 1
|
bestPartialLen[hunkIndex] = i - 1
|
||||||
bestPartialAnchorLine[hunkIndex] = lineIndex
|
bestPartialAnchorLine[hunkIndex] = lineIndex
|
||||||
@@ -646,9 +716,13 @@ patch_file() {
|
|||||||
print "" > "/dev/stderr"
|
print "" > "/dev/stderr"
|
||||||
print "Closest match: anchored at file line " bestPartialAnchorLine[failingHunk] ", matched " bestPartialLen[failingHunk] " of " hunkTotalOriginalLines[failingHunk] " original lines before diverging." > "/dev/stderr"
|
print "Closest match: anchored at file line " bestPartialAnchorLine[failingHunk] ", matched " bestPartialLen[failingHunk] " of " hunkTotalOriginalLines[failingHunk] " original lines before diverging." > "/dev/stderr"
|
||||||
print "" > "/dev/stderr"
|
print "" > "/dev/stderr"
|
||||||
print "At file line " bestPartialDivergeLine[failingHunk] " (hunk original line " bestPartialHunkPos[failingHunk] "):" > "/dev/stderr"
|
if (bestPartialDivergeLine[failingHunk] > totalLines) {
|
||||||
print " expected: " bestPartialExpected[failingHunk] > "/dev/stderr"
|
print "The hunk expects additional lines beyond the end of the file (file has " totalLines " lines)." > "/dev/stderr"
|
||||||
print " actual: " bestPartialActual[failingHunk] > "/dev/stderr"
|
} else {
|
||||||
|
print "At file line " bestPartialDivergeLine[failingHunk] " (hunk original line " bestPartialHunkPos[failingHunk] "):" > "/dev/stderr"
|
||||||
|
print " expected: " bestPartialExpected[failingHunk] > "/dev/stderr"
|
||||||
|
print " actual: " bestPartialActual[failingHunk] > "/dev/stderr"
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
print "" > "/dev/stderr"
|
print "" > "/dev/stderr"
|
||||||
|
|||||||
@@ -2,7 +2,7 @@
|
|||||||
description: Navigate and curate markdown knowledge bases (plan repos, spec repos, companion docs) with IWE graph tools. Load when the workspace is or contains a markdown knowledge base and the task involves finding, reading, or reorganizing plans, specs, designs, or notes. Activates the iwe MCP server rooted at the current directory.
|
description: Navigate and curate markdown knowledge bases (plan repos, spec repos, companion docs) with IWE graph tools. Load when the workspace is or contains a markdown knowledge base and the task involves finding, reading, or reorganizing plans, specs, designs, or notes. Activates the iwe MCP server rooted at the current directory.
|
||||||
enabled_mcp_servers: iwe
|
enabled_mcp_servers: iwe
|
||||||
---
|
---
|
||||||
You are working with a markdown knowledge base through IWE, a graph-based knowledge tool. The `iwe` MCP server is rooted at the current working directory (`--project .`), so the knowledge base is the directory Coyote was launched in. IWE derives structure from links: a link on its own line is an *inclusion link* (parent-child hierarchy); a link inside text is an *inline reference* (cross-reference, produces backlinks). The server watches the filesystem, so external edits are picked up automatically — never ask for a restart.
|
You are working with a markdown knowledge base through IWE, a graph-based knowledge tool. The `iwe` MCP server is rooted at the current working directory, so the knowledge base is the directory Coyote was launched in. IWE derives structure from links: a link on its own line is an *inclusion link* (parent-child hierarchy); a link inside text is an *inline reference* (cross-reference, produces backlinks). The server watches the filesystem, so external edits are picked up automatically — never ask for a restart.
|
||||||
|
|
||||||
## When to use this (and when not)
|
## When to use this (and when not)
|
||||||
|
|
||||||
|
|||||||
+2
-1
@@ -292,6 +292,7 @@ clients:
|
|||||||
# extra:
|
# extra:
|
||||||
# proxy: socks5://127.0.0.1:1080 # Set proxy
|
# proxy: socks5://127.0.0.1:1080 # Set proxy
|
||||||
# connect_timeout: 10 # Set timeout in seconds for connect to api
|
# connect_timeout: 10 # Set timeout in seconds for connect to api
|
||||||
|
# read_timeout: 300 # Set timeout in seconds for a read stall (no bytes received); 0 disables (default: 300)
|
||||||
|
|
||||||
# See https://platform.openai.com/docs/quickstart
|
# See https://platform.openai.com/docs/quickstart
|
||||||
- type: openai
|
- type: openai
|
||||||
@@ -525,4 +526,4 @@ clients:
|
|||||||
- type: openai-compatible
|
- type: openai-compatible
|
||||||
name: voyageai
|
name: voyageai
|
||||||
api_base: https://api.voyageai.com/v1
|
api_base: https://api.voyageai.com/v1
|
||||||
api_key: '{{VOYAGEAI_API_KEY}}' # You can either hard-code or inject secrets from the Coyote vault
|
api_key: '{{VOYAGEAI_API_KEY}}' # You can either hard-code or inject secrets from the Coyote vault
|
||||||
@@ -841,6 +841,11 @@
|
|||||||
referrer: coyote
|
referrer: coyote
|
||||||
echo_pkce_in_token_exchange: true
|
echo_pkce_in_token_exchange: true
|
||||||
models:
|
models:
|
||||||
|
- name: grok-4.6
|
||||||
|
input_price: 2
|
||||||
|
output_price: 6
|
||||||
|
max_input_tokens: 500000
|
||||||
|
supports_function_calling: true
|
||||||
- name: grok-4.5
|
- name: grok-4.5
|
||||||
input_price: 2
|
input_price: 2
|
||||||
output_price: 6
|
output_price: 6
|
||||||
|
|||||||
+116
-1
@@ -2,6 +2,7 @@ use anyhow::{Result, anyhow};
|
|||||||
use chrono::Utc;
|
use chrono::Utc;
|
||||||
use indexmap::IndexMap;
|
use indexmap::IndexMap;
|
||||||
use parking_lot::RwLock;
|
use parking_lot::RwLock;
|
||||||
|
use std::collections::HashMap;
|
||||||
use std::sync::LazyLock;
|
use std::sync::LazyLock;
|
||||||
|
|
||||||
type AccessTokenEntry = (String, i64, Option<String>);
|
type AccessTokenEntry = (String, i64, Option<String>);
|
||||||
@@ -9,6 +10,12 @@ type AccessTokenEntry = (String, i64, Option<String>);
|
|||||||
static ACCESS_TOKENS: LazyLock<RwLock<IndexMap<String, AccessTokenEntry>>> =
|
static ACCESS_TOKENS: LazyLock<RwLock<IndexMap<String, AccessTokenEntry>>> =
|
||||||
LazyLock::new(|| RwLock::new(IndexMap::new()));
|
LazyLock::new(|| RwLock::new(IndexMap::new()));
|
||||||
|
|
||||||
|
/// Tokens a provider rejected (401) despite being locally unexpired.
|
||||||
|
/// Maps client name → the exact rejected token so a concurrently-refreshed
|
||||||
|
/// different token is never distrusted by mistake.
|
||||||
|
static REJECTED_TOKENS: LazyLock<RwLock<HashMap<String, String>>> =
|
||||||
|
LazyLock::new(|| RwLock::new(HashMap::new()));
|
||||||
|
|
||||||
pub fn get_access_token(client_name: &str) -> Result<String> {
|
pub fn get_access_token(client_name: &str) -> Result<String> {
|
||||||
ACCESS_TOKENS
|
ACCESS_TOKENS
|
||||||
.read()
|
.read()
|
||||||
@@ -30,7 +37,7 @@ pub fn is_valid_access_token(client_name: &str) -> bool {
|
|||||||
Some(v) => v,
|
Some(v) => v,
|
||||||
None => return false,
|
None => return false,
|
||||||
};
|
};
|
||||||
!token.is_empty() && Utc::now().timestamp() < *expires_at
|
!token.is_empty() && Utc::now().timestamp() < *expires_at && !is_rejected(client_name, token)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn set_access_token(
|
pub fn set_access_token(
|
||||||
@@ -45,3 +52,111 @@ pub fn set_access_token(
|
|||||||
entry.1 = expires_at;
|
entry.1 = expires_at;
|
||||||
entry.2 = account_id;
|
entry.2 = account_id;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Compare-and-invalidate a provider-rejected token.
|
||||||
|
///
|
||||||
|
/// Only if the currently-cached token EQUALS `rejected` is the cache entry
|
||||||
|
/// removed and the rejection marker recorded; a concurrently-refreshed
|
||||||
|
/// different token is left untouched and no marker is set.
|
||||||
|
///
|
||||||
|
/// Returns true if a cache entry existed for this client at all (whether or
|
||||||
|
/// not it matched `rejected`) — i.e. the client is token-authed and a retry
|
||||||
|
/// after refresh is worthwhile. Returns false when there is no entry
|
||||||
|
/// (API-key clients).
|
||||||
|
pub fn distrust_access_token(client_name: &str, rejected: &str) -> bool {
|
||||||
|
let mut access_tokens = ACCESS_TOKENS.write();
|
||||||
|
let (token, _, _) = match access_tokens.get(client_name) {
|
||||||
|
Some(v) => v,
|
||||||
|
None => return false,
|
||||||
|
};
|
||||||
|
if token == rejected {
|
||||||
|
access_tokens.shift_remove(client_name);
|
||||||
|
REJECTED_TOKENS
|
||||||
|
.write()
|
||||||
|
.insert(client_name.to_string(), rejected.to_string());
|
||||||
|
}
|
||||||
|
true
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn is_rejected(client_name: &str, token: &str) -> bool {
|
||||||
|
REJECTED_TOKENS
|
||||||
|
.read()
|
||||||
|
.get(client_name)
|
||||||
|
.is_some_and(|rejected| rejected == token)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn clear_rejected(client_name: &str) {
|
||||||
|
REJECTED_TOKENS.write().remove(client_name);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn distrust_removes_matching_token_and_sets_marker() {
|
||||||
|
let client = "distrust-match-test";
|
||||||
|
set_access_token(client, "at-1".into(), Utc::now().timestamp() + 3600, None);
|
||||||
|
|
||||||
|
assert!(distrust_access_token(client, "at-1"));
|
||||||
|
|
||||||
|
assert!(get_access_token(client).is_err(), "cache entry not removed");
|
||||||
|
assert!(is_rejected(client, "at-1"), "marker not set");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn distrust_keeps_differing_token_and_skips_marker() {
|
||||||
|
let client = "distrust-differ-test";
|
||||||
|
set_access_token(client, "at-new".into(), Utc::now().timestamp() + 3600, None);
|
||||||
|
|
||||||
|
assert!(distrust_access_token(client, "at-old"));
|
||||||
|
|
||||||
|
assert_eq!(get_access_token(client).unwrap(), "at-new");
|
||||||
|
assert!(!is_rejected(client, "at-old"), "marker set for stale token");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn distrust_returns_false_without_cache_entry() {
|
||||||
|
let client = "distrust-missing-test";
|
||||||
|
|
||||||
|
assert!(!distrust_access_token(client, "at-1"));
|
||||||
|
assert!(!is_rejected(client, "at-1"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn is_valid_access_token_false_for_rejected_token() {
|
||||||
|
let client = "rejected-valid-test";
|
||||||
|
set_access_token(client, "at-1".into(), Utc::now().timestamp() + 3600, None);
|
||||||
|
assert!(is_valid_access_token(client));
|
||||||
|
|
||||||
|
distrust_access_token(client, "at-1");
|
||||||
|
// A concurrent in-flight prepare re-caches the rejected file token
|
||||||
|
// between mark and refresh; it must still be treated as invalid.
|
||||||
|
set_access_token(client, "at-1".into(), Utc::now().timestamp() + 3600, None);
|
||||||
|
|
||||||
|
assert!(!is_valid_access_token(client));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn clear_rejected_clears_marker_and_clients_are_isolated() {
|
||||||
|
let client_a = "rejected-isolation-a";
|
||||||
|
let client_b = "rejected-isolation-b";
|
||||||
|
set_access_token(client_a, "at-1".into(), Utc::now().timestamp() + 3600, None);
|
||||||
|
distrust_access_token(client_a, "at-1");
|
||||||
|
|
||||||
|
assert!(is_rejected(client_a, "at-1"));
|
||||||
|
assert!(
|
||||||
|
!is_rejected(client_b, "at-1"),
|
||||||
|
"marker leaked across clients"
|
||||||
|
);
|
||||||
|
|
||||||
|
clear_rejected(client_b);
|
||||||
|
assert!(
|
||||||
|
is_rejected(client_a, "at-1"),
|
||||||
|
"wrong client's marker cleared"
|
||||||
|
);
|
||||||
|
|
||||||
|
clear_rejected(client_a);
|
||||||
|
assert!(!is_rejected(client_a, "at-1"));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
+289
-15
@@ -1,5 +1,6 @@
|
|||||||
use super::*;
|
use super::*;
|
||||||
|
|
||||||
|
use super::access_token::{distrust_access_token, get_access_token};
|
||||||
use crate::config::{RenderMode, paths};
|
use crate::config::{RenderMode, paths};
|
||||||
use crate::{
|
use crate::{
|
||||||
config::{AppConfig, Input, RequestContext},
|
config::{AppConfig, Input, RequestContext},
|
||||||
@@ -56,12 +57,16 @@ pub trait Client: Sync + Send {
|
|||||||
let mut builder = ReqwestClient::builder();
|
let mut builder = ReqwestClient::builder();
|
||||||
let extra = self.extra_config();
|
let extra = self.extra_config();
|
||||||
let timeout = extra.and_then(|v| v.connect_timeout).unwrap_or(10);
|
let timeout = extra.and_then(|v| v.connect_timeout).unwrap_or(10);
|
||||||
|
let read_timeout = extra.and_then(|v| v.read_timeout).unwrap_or(300);
|
||||||
if let Some(proxy) = extra.and_then(|v| v.proxy.as_deref()) {
|
if let Some(proxy) = extra.and_then(|v| v.proxy.as_deref()) {
|
||||||
builder = set_proxy(builder, proxy)?;
|
builder = set_proxy(builder, proxy)?;
|
||||||
}
|
}
|
||||||
if let Some(user_agent) = self.app_config().user_agent.as_ref() {
|
if let Some(user_agent) = self.app_config().user_agent.as_ref() {
|
||||||
builder = builder.user_agent(user_agent);
|
builder = builder.user_agent(user_agent);
|
||||||
}
|
}
|
||||||
|
if read_timeout > 0 {
|
||||||
|
builder = builder.read_timeout(Duration::from_secs(read_timeout));
|
||||||
|
}
|
||||||
let client = builder
|
let client = builder
|
||||||
.connect_timeout(Duration::from_secs(timeout))
|
.connect_timeout(Duration::from_secs(timeout))
|
||||||
.build()
|
.build()
|
||||||
@@ -69,6 +74,11 @@ pub trait Client: Sync + Send {
|
|||||||
Ok(client)
|
Ok(client)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// On a 401 the cached access token is distrusted and the call retried
|
||||||
|
/// exactly once; the retry re-runs the per-client prepare step, which
|
||||||
|
/// sees the rejection marker, force-refreshes the token, and rebuilds
|
||||||
|
/// the whole request. A second 401 propagates the original error; any
|
||||||
|
/// other retry failure propagates as-is.
|
||||||
async fn chat_completions(&self, input: Input) -> Result<ChatCompletionsOutput> {
|
async fn chat_completions(&self, input: Input) -> Result<ChatCompletionsOutput> {
|
||||||
if self.app_config().dry_run {
|
if self.app_config().dry_run {
|
||||||
let content = input.echo_messages();
|
let content = input.echo_messages();
|
||||||
@@ -76,11 +86,30 @@ pub trait Client: Sync + Send {
|
|||||||
}
|
}
|
||||||
let client = self.build_client()?;
|
let client = self.build_client()?;
|
||||||
let data = input.prepare_completion_data(self.model(), false)?;
|
let data = input.prepare_completion_data(self.model(), false)?;
|
||||||
self.chat_completions_inner(&client, data)
|
let err = match self.chat_completions_inner(&client, data).await {
|
||||||
.await
|
Ok(output) => return Ok(output),
|
||||||
.with_context(|| "Failed to call chat-completions api")
|
Err(err) => err,
|
||||||
|
};
|
||||||
|
let ret = if should_retry_auth(&err, self.name()) {
|
||||||
|
debug!(
|
||||||
|
"provider '{}' rejected access token (401); refreshing and retrying once",
|
||||||
|
self.name()
|
||||||
|
);
|
||||||
|
let data = input.prepare_completion_data(self.model(), false)?;
|
||||||
|
match self.chat_completions_inner(&client, data).await {
|
||||||
|
Err(retry_err) if is_auth_error(&retry_err) => Err(err),
|
||||||
|
ret => ret,
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
Err(err)
|
||||||
|
};
|
||||||
|
ret.with_context(|| "Failed to call chat-completions api")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Same retry-once-on-401 semantics as [`Self::chat_completions`], but
|
||||||
|
/// only while the handler has received nothing yet: retrying after
|
||||||
|
/// partial output has streamed would render it to the user twice. The
|
||||||
|
/// retry lives inside the same `select!` arm so abort stays responsive.
|
||||||
async fn chat_completions_streaming(
|
async fn chat_completions_streaming(
|
||||||
&self,
|
&self,
|
||||||
input: &Input,
|
input: &Input,
|
||||||
@@ -97,7 +126,22 @@ pub trait Client: Sync + Send {
|
|||||||
}
|
}
|
||||||
let client = self.build_client()?;
|
let client = self.build_client()?;
|
||||||
let data = input.prepare_completion_data(self.model(), true)?;
|
let data = input.prepare_completion_data(self.model(), true)?;
|
||||||
self.chat_completions_streaming_inner(&client, handler, data).await
|
let err = match self.chat_completions_streaming_inner(&client, handler, data).await {
|
||||||
|
Ok(()) => return Ok(()),
|
||||||
|
Err(err) => err,
|
||||||
|
};
|
||||||
|
if handler.has_received_content() || !should_retry_auth(&err, self.name()) {
|
||||||
|
return Err(err);
|
||||||
|
}
|
||||||
|
debug!(
|
||||||
|
"provider '{}' rejected access token (401); refreshing and retrying once",
|
||||||
|
self.name()
|
||||||
|
);
|
||||||
|
let data = input.prepare_completion_data(self.model(), true)?;
|
||||||
|
match self.chat_completions_streaming_inner(&client, handler, data).await {
|
||||||
|
Err(retry_err) if is_auth_error(&retry_err) => Err(err),
|
||||||
|
ret => ret,
|
||||||
|
}
|
||||||
} => {
|
} => {
|
||||||
handler.done();
|
handler.done();
|
||||||
ret.with_context(|| "Failed to call chat-completions api")
|
ret.with_context(|| "Failed to call chat-completions api")
|
||||||
@@ -109,11 +153,27 @@ pub trait Client: Sync + Send {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Same retry-once-on-401 semantics as [`Self::chat_completions`]
|
||||||
|
/// (gemini OAuth embeddings route here).
|
||||||
async fn embeddings(&self, data: &EmbeddingsData) -> Result<Vec<Vec<f32>>> {
|
async fn embeddings(&self, data: &EmbeddingsData) -> Result<Vec<Vec<f32>>> {
|
||||||
let client = self.build_client()?;
|
let client = self.build_client()?;
|
||||||
self.embeddings_inner(&client, data)
|
let err = match self.embeddings_inner(&client, data).await {
|
||||||
.await
|
Ok(output) => return Ok(output),
|
||||||
.context("Failed to call embeddings api")
|
Err(err) => err,
|
||||||
|
};
|
||||||
|
let ret = if should_retry_auth(&err, self.name()) {
|
||||||
|
debug!(
|
||||||
|
"provider '{}' rejected access token (401); refreshing and retrying once",
|
||||||
|
self.name()
|
||||||
|
);
|
||||||
|
match self.embeddings_inner(&client, data).await {
|
||||||
|
Err(retry_err) if is_auth_error(&retry_err) => Err(err),
|
||||||
|
ret => ret,
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
Err(err)
|
||||||
|
};
|
||||||
|
ret.context("Failed to call embeddings api")
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn rerank(&self, data: &RerankData) -> Result<RerankOutput> {
|
async fn rerank(&self, data: &RerankData) -> Result<RerankOutput> {
|
||||||
@@ -205,6 +265,7 @@ impl Default for ClientConfig {
|
|||||||
pub struct ExtraConfig {
|
pub struct ExtraConfig {
|
||||||
pub proxy: Option<String>,
|
pub proxy: Option<String>,
|
||||||
pub connect_timeout: Option<u64>,
|
pub connect_timeout: Option<u64>,
|
||||||
|
pub read_timeout: Option<u64>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Deserialize, Default)]
|
#[derive(Debug, Clone, Deserialize, Default)]
|
||||||
@@ -557,46 +618,90 @@ pub async fn noop_rerank(_builder: RequestBuilder, _model: &Model) -> Result<Rer
|
|||||||
bail!("The client doesn't support rerank api")
|
bail!("The client doesn't support rerank api")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub struct ApiStatusError {
|
||||||
|
pub status: u16,
|
||||||
|
pub message: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl std::fmt::Display for ApiStatusError {
|
||||||
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||||
|
write!(f, "{}", self.message)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl std::error::Error for ApiStatusError {}
|
||||||
|
|
||||||
|
/// True when the error chain bottoms out in an [`ApiStatusError`] with
|
||||||
|
/// status 401 EXACTLY. 403 (entitlement) and 429 (rate limit) are never
|
||||||
|
/// auth failures, and message text is never inspected.
|
||||||
|
fn is_auth_error(err: &anyhow::Error) -> bool {
|
||||||
|
err.downcast_ref::<ApiStatusError>()
|
||||||
|
.is_some_and(|api_err| api_err.status == 401)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Decides whether a 401 from `client_name` warrants a single retry after a
|
||||||
|
/// forced token refresh: the error must be a 401 [`ApiStatusError`], and the
|
||||||
|
/// client must have a cached access token to distrust (API-key clients have
|
||||||
|
/// none and never retry). Distrusting marks the exact rejected token so the
|
||||||
|
/// retry's prepare step force-refreshes it. There is deliberately no backoff:
|
||||||
|
/// the blast radius is bounded at one extra request per user-visible call.
|
||||||
|
///
|
||||||
|
/// Note: vertexai shares the ACCESS_TOKENS cache, so a 401 there also
|
||||||
|
/// triggers distrust+retry — deliberate.
|
||||||
|
fn should_retry_auth(err: &anyhow::Error, client_name: &str) -> bool {
|
||||||
|
if !is_auth_error(err) {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
let Ok(token) = get_access_token(client_name) else {
|
||||||
|
return false;
|
||||||
|
};
|
||||||
|
distrust_access_token(client_name, &token)
|
||||||
|
}
|
||||||
|
|
||||||
pub fn catch_error(data: &Value, status: u16) -> Result<()> {
|
pub fn catch_error(data: &Value, status: u16) -> Result<()> {
|
||||||
if (200..300).contains(&status) {
|
if (200..300).contains(&status) {
|
||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
debug!("Invalid response, status: {status}, data: {data}");
|
debug!("Invalid response, status: {status}, data: {data}");
|
||||||
|
let api_error = |message: String| anyhow::Error::new(ApiStatusError { status, message });
|
||||||
if let Some(error) = data["error"].as_object() {
|
if let Some(error) = data["error"].as_object() {
|
||||||
if let (Some(typ), Some(message)) = (
|
if let (Some(typ), Some(message)) = (
|
||||||
json_str_from_map(error, "type"),
|
json_str_from_map(error, "type"),
|
||||||
json_str_from_map(error, "message"),
|
json_str_from_map(error, "message"),
|
||||||
) {
|
) {
|
||||||
bail!("{message} (type: {typ})");
|
return Err(api_error(format!("{message} (type: {typ})")));
|
||||||
} else if let (Some(typ), Some(message)) = (
|
} else if let (Some(typ), Some(message)) = (
|
||||||
json_str_from_map(error, "code"),
|
json_str_from_map(error, "code"),
|
||||||
json_str_from_map(error, "message"),
|
json_str_from_map(error, "message"),
|
||||||
) {
|
) {
|
||||||
bail!("{message} (code: {typ})");
|
return Err(api_error(format!("{message} (code: {typ})")));
|
||||||
}
|
}
|
||||||
} else if let Some(error) = data["errors"][0].as_object() {
|
} else if let Some(error) = data["errors"][0].as_object() {
|
||||||
if let (Some(code), Some(message)) = (
|
if let (Some(code), Some(message)) = (
|
||||||
error.get("code").and_then(|v| v.as_u64()),
|
error.get("code").and_then(|v| v.as_u64()),
|
||||||
json_str_from_map(error, "message"),
|
json_str_from_map(error, "message"),
|
||||||
) {
|
) {
|
||||||
bail!("{message} (status: {code})")
|
return Err(api_error(format!("{message} (status: {code})")));
|
||||||
}
|
}
|
||||||
} else if let Some(error) = data[0]["error"].as_object() {
|
} else if let Some(error) = data[0]["error"].as_object() {
|
||||||
if let (Some(status), Some(message)) = (
|
if let (Some(status), Some(message)) = (
|
||||||
json_str_from_map(error, "status"),
|
json_str_from_map(error, "status"),
|
||||||
json_str_from_map(error, "message"),
|
json_str_from_map(error, "message"),
|
||||||
) {
|
) {
|
||||||
bail!("{message} (status: {status})")
|
return Err(api_error(format!("{message} (status: {status})")));
|
||||||
}
|
}
|
||||||
} else if let (Some(detail), Some(status)) = (data["detail"].as_str(), data["status"].as_i64())
|
} else if let (Some(detail), Some(status)) = (data["detail"].as_str(), data["status"].as_i64())
|
||||||
{
|
{
|
||||||
bail!("{detail} (status: {status})");
|
return Err(api_error(format!("{detail} (status: {status})")));
|
||||||
} else if let Some(error) = data["error"].as_str() {
|
} else if let Some(error) = data["error"].as_str() {
|
||||||
bail!("{error}");
|
return Err(api_error(error.to_string()));
|
||||||
} else if let Some(message) = data["message"].as_str() {
|
} else if let Some(message) = data["message"].as_str() {
|
||||||
bail!("{message}");
|
return Err(api_error(message.to_string()));
|
||||||
}
|
}
|
||||||
bail!("Invalid response data: {data} (status: {status})");
|
Err(api_error(format!(
|
||||||
|
"Invalid response data: {data} (status: {status})"
|
||||||
|
)))
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn json_str_from_map<'a>(
|
pub fn json_str_from_map<'a>(
|
||||||
@@ -737,3 +842,172 @@ fn prompt_input_string(desc: &str, required: bool, help_message: Option<&str>) -
|
|||||||
let text = text.prompt()?;
|
let text = text.prompt()?;
|
||||||
Ok(text)
|
Ok(text)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
use super::super::access_token::{is_rejected, set_access_token};
|
||||||
|
|
||||||
|
fn catch_error_message(data: &Value, status: u16) -> String {
|
||||||
|
catch_error(data, status).unwrap_err().to_string()
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_catch_error_display_json_with_type() {
|
||||||
|
let data = json!({"error": {"type": "invalid_request_error", "message": "Bad request"}});
|
||||||
|
assert_eq!(
|
||||||
|
catch_error_message(&data, 400),
|
||||||
|
"Bad request (type: invalid_request_error)"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_catch_error_display_json_with_code() {
|
||||||
|
let data = json!({"error": {"code": "rate_limited", "message": "Too many requests"}});
|
||||||
|
assert_eq!(
|
||||||
|
catch_error_message(&data, 429),
|
||||||
|
"Too many requests (code: rate_limited)"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_catch_error_display_errors_array() {
|
||||||
|
let data = json!({"errors": [{"code": 7000, "message": "No route"}]});
|
||||||
|
assert_eq!(catch_error_message(&data, 404), "No route (status: 7000)");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_catch_error_display_array_error_status() {
|
||||||
|
let data = json!([{"error": {"status": "PERMISSION_DENIED", "message": "Denied"}}]);
|
||||||
|
assert_eq!(
|
||||||
|
catch_error_message(&data, 403),
|
||||||
|
"Denied (status: PERMISSION_DENIED)"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_catch_error_display_detail_status() {
|
||||||
|
let data = json!({"detail": "Not found", "status": 404});
|
||||||
|
assert_eq!(catch_error_message(&data, 404), "Not found (status: 404)");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_catch_error_display_error_string() {
|
||||||
|
let data = json!({"error": "Something went wrong"});
|
||||||
|
assert_eq!(catch_error_message(&data, 500), "Something went wrong");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_catch_error_display_message_string() {
|
||||||
|
let data = json!({"message": "Unauthorized"});
|
||||||
|
assert_eq!(catch_error_message(&data, 401), "Unauthorized");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_catch_error_display_fallback() {
|
||||||
|
let data = json!({"unexpected": true});
|
||||||
|
assert_eq!(
|
||||||
|
catch_error_message(&data, 500),
|
||||||
|
format!("Invalid response data: {data} (status: 500)")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_catch_error_ok_on_success_status() {
|
||||||
|
let data = json!({"error": {"type": "x", "message": "y"}});
|
||||||
|
assert!(catch_error(&data, 200).is_ok());
|
||||||
|
assert!(catch_error(&data, 299).is_ok());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_catch_error_downcast_through_context_chain() {
|
||||||
|
let data = json!({"error": {"type": "authentication_error", "message": "Invalid key"}});
|
||||||
|
let err = catch_error(&data, 401)
|
||||||
|
.context("Failed to call chat-completions api")
|
||||||
|
.unwrap_err();
|
||||||
|
let api_err = err
|
||||||
|
.downcast_ref::<ApiStatusError>()
|
||||||
|
.expect("should downcast through context chain");
|
||||||
|
assert_eq!(api_err.status, 401);
|
||||||
|
assert_eq!(api_err.message, "Invalid key (type: authentication_error)");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_catch_error_preserves_status() {
|
||||||
|
let data = json!({"message": "Unauthorized"});
|
||||||
|
let err = catch_error(&data, 401).unwrap_err();
|
||||||
|
assert_eq!(err.downcast_ref::<ApiStatusError>().unwrap().status, 401);
|
||||||
|
|
||||||
|
let data = json!({"detail": "Rate limited", "status": 429});
|
||||||
|
let err = catch_error(&data, 429).unwrap_err();
|
||||||
|
assert_eq!(err.downcast_ref::<ApiStatusError>().unwrap().status, 429);
|
||||||
|
|
||||||
|
// The struct carries the outer HTTP status even when the body embeds another code
|
||||||
|
let data = json!({"errors": [{"code": 7000, "message": "No route"}]});
|
||||||
|
let err = catch_error(&data, 429).unwrap_err();
|
||||||
|
assert_eq!(err.downcast_ref::<ApiStatusError>().unwrap().status, 429);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Wrapped in `.context(...)` so every test below proves the downcast
|
||||||
|
/// works through an anyhow context chain, as in the trait methods.
|
||||||
|
fn api_status_error(status: u16) -> anyhow::Error {
|
||||||
|
anyhow::Error::new(ApiStatusError {
|
||||||
|
status,
|
||||||
|
message: format!("error (status: {status})"),
|
||||||
|
})
|
||||||
|
.context("Failed to call chat-completions api")
|
||||||
|
}
|
||||||
|
|
||||||
|
fn cache_token(client: &str, token: &str) {
|
||||||
|
set_access_token(
|
||||||
|
client,
|
||||||
|
token.into(),
|
||||||
|
chrono::Utc::now().timestamp() + 3600,
|
||||||
|
None,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_should_retry_auth_401_with_cached_token() {
|
||||||
|
let client = "should-retry-auth-401";
|
||||||
|
cache_token(client, "at-1");
|
||||||
|
|
||||||
|
assert!(should_retry_auth(&api_status_error(401), client));
|
||||||
|
assert!(is_rejected(client, "at-1"), "rejected marker not set");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_should_retry_auth_non_401_statuses() {
|
||||||
|
let client = "should-retry-auth-non-401";
|
||||||
|
cache_token(client, "at-1");
|
||||||
|
|
||||||
|
for status in [403, 429, 500] {
|
||||||
|
assert!(
|
||||||
|
!should_retry_auth(&api_status_error(status), client),
|
||||||
|
"retried on {status}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
assert_eq!(get_access_token(client).unwrap(), "at-1");
|
||||||
|
assert!(!is_rejected(client, "at-1"), "marker set without a 401");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_should_retry_auth_non_api_status_error() {
|
||||||
|
let client = "should-retry-auth-non-api";
|
||||||
|
cache_token(client, "at-1");
|
||||||
|
|
||||||
|
let err = anyhow::anyhow!("connection reset").context("Failed to call embeddings api");
|
||||||
|
assert!(!should_retry_auth(&err, client));
|
||||||
|
assert_eq!(get_access_token(client).unwrap(), "at-1");
|
||||||
|
assert!(!is_rejected(client, "at-1"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_should_retry_auth_401_without_cached_token() {
|
||||||
|
let client = "should-retry-auth-no-token";
|
||||||
|
|
||||||
|
assert!(!should_retry_auth(&api_status_error(401), client));
|
||||||
|
assert!(!is_rejected(client, "at-1"));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
+410
-43
@@ -1,14 +1,14 @@
|
|||||||
use super::access_token::{is_valid_access_token, set_access_token};
|
use super::access_token::{clear_rejected, is_rejected, is_valid_access_token, set_access_token};
|
||||||
use super::openai_compatible_oauth::OpenAICompatibleOAuthProvider;
|
use super::openai_compatible_oauth::OpenAICompatibleOAuthProvider;
|
||||||
use super::{ClientConfig, ProviderModels};
|
use super::{ClientConfig, ProviderModels};
|
||||||
use crate::config::paths;
|
use crate::config::paths;
|
||||||
use anyhow::{Context, Result, anyhow, bail};
|
use anyhow::{Context, Error, Result, anyhow, bail};
|
||||||
use base64::Engine;
|
use base64::Engine;
|
||||||
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
|
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
|
||||||
use chrono::Utc;
|
use chrono::Utc;
|
||||||
use indexmap::IndexMap;
|
use indexmap::IndexMap;
|
||||||
use inquire::Text;
|
use inquire::Text;
|
||||||
use reqwest::{Client as ReqwestClient, RequestBuilder};
|
use reqwest::{Client as ReqwestClient, RequestBuilder, StatusCode};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use serde_json::Value;
|
use serde_json::Value;
|
||||||
use sha2::{Digest, Sha256};
|
use sha2::{Digest, Sha256};
|
||||||
@@ -16,6 +16,10 @@ use std::collections::HashMap;
|
|||||||
use std::fs;
|
use std::fs;
|
||||||
use std::io::{BufRead, BufReader, Write};
|
use std::io::{BufRead, BufReader, Write};
|
||||||
use std::net::TcpListener;
|
use std::net::TcpListener;
|
||||||
|
use std::path::PathBuf;
|
||||||
|
use std::sync::{Arc, OnceLock};
|
||||||
|
use std::time::Duration;
|
||||||
|
use tokio::sync;
|
||||||
use url::Url;
|
use url::Url;
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
@@ -197,10 +201,20 @@ pub struct OAuthTokens {
|
|||||||
pub account_id: Option<String>,
|
pub account_id: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
const TOKEN_ENDPOINT_TIMEOUT: Duration = Duration::from_secs(30);
|
||||||
|
|
||||||
pub async fn run_oauth_flow(provider: &dyn OAuthProvider, client_name: &str) -> Result<()> {
|
pub async fn run_oauth_flow(provider: &dyn OAuthProvider, client_name: &str) -> Result<()> {
|
||||||
match provider.flow() {
|
match provider.flow() {
|
||||||
OAuthFlow::Pkce => run_pkce_flow(provider, client_name).await,
|
OAuthFlow::Pkce => run_pkce_flow(provider, client_name).await,
|
||||||
OAuthFlow::ClientCredentials => run_client_credentials_flow(provider, client_name).await,
|
OAuthFlow::ClientCredentials => {
|
||||||
|
run_client_credentials_flow(provider, client_name).await?;
|
||||||
|
println!(
|
||||||
|
"Successfully authenticated client '{}' with {} via OAuth (client_credentials). Tokens saved.",
|
||||||
|
client_name,
|
||||||
|
provider.provider_name()
|
||||||
|
);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
OAuthFlow::DeviceCode => run_device_code_flow(provider, client_name).await,
|
OAuthFlow::DeviceCode => run_device_code_flow(provider, client_name).await,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -301,12 +315,20 @@ async fn run_pkce_flow(provider: &dyn OAuthProvider, client_name: &str) -> Resul
|
|||||||
|
|
||||||
let access_token = response["access_token"]
|
let access_token = response["access_token"]
|
||||||
.as_str()
|
.as_str()
|
||||||
.ok_or_else(|| anyhow!("Missing access_token in response: {response}"))?
|
.ok_or_else(|| {
|
||||||
|
anyhow!(
|
||||||
|
"Missing access_token in response (keys: {})",
|
||||||
|
token_response_keys(&response)
|
||||||
|
)
|
||||||
|
})?
|
||||||
.to_string();
|
.to_string();
|
||||||
let refresh_token = response["refresh_token"].as_str().map(|s| s.to_string());
|
let refresh_token = response["refresh_token"].as_str().map(|s| s.to_string());
|
||||||
let expires_in = response["expires_in"]
|
let expires_in = response["expires_in"].as_i64().ok_or_else(|| {
|
||||||
.as_i64()
|
anyhow!(
|
||||||
.ok_or_else(|| anyhow!("Missing expires_in in response: {response}"))?;
|
"Missing expires_in in response (keys: {})",
|
||||||
|
token_response_keys(&response)
|
||||||
|
)
|
||||||
|
})?;
|
||||||
|
|
||||||
let expires_at = Utc::now().timestamp() + expires_in;
|
let expires_at = Utc::now().timestamp() + expires_in;
|
||||||
|
|
||||||
@@ -334,7 +356,9 @@ async fn run_client_credentials_flow(
|
|||||||
provider: &dyn OAuthProvider,
|
provider: &dyn OAuthProvider,
|
||||||
client_name: &str,
|
client_name: &str,
|
||||||
) -> Result<()> {
|
) -> Result<()> {
|
||||||
let client = ReqwestClient::new();
|
let client = ReqwestClient::builder()
|
||||||
|
.timeout(TOKEN_ENDPOINT_TIMEOUT)
|
||||||
|
.build()?;
|
||||||
let scopes = provider.scopes();
|
let scopes = provider.scopes();
|
||||||
let mut params: Vec<(&str, &str)> = vec![
|
let mut params: Vec<(&str, &str)> = vec![
|
||||||
("grant_type", "client_credentials"),
|
("grant_type", "client_credentials"),
|
||||||
@@ -349,11 +373,19 @@ async fn run_client_credentials_flow(
|
|||||||
|
|
||||||
let access_token = response["access_token"]
|
let access_token = response["access_token"]
|
||||||
.as_str()
|
.as_str()
|
||||||
.ok_or_else(|| anyhow!("Missing access_token in client_credentials response: {response}"))?
|
.ok_or_else(|| {
|
||||||
|
anyhow!(
|
||||||
|
"Missing access_token in client_credentials response (keys: {})",
|
||||||
|
token_response_keys(&response)
|
||||||
|
)
|
||||||
|
})?
|
||||||
.to_string();
|
.to_string();
|
||||||
let expires_in = response["expires_in"]
|
let expires_in = response["expires_in"].as_i64().ok_or_else(|| {
|
||||||
.as_i64()
|
anyhow!(
|
||||||
.ok_or_else(|| anyhow!("Missing expires_in in client_credentials response: {response}"))?;
|
"Missing expires_in in client_credentials response (keys: {})",
|
||||||
|
token_response_keys(&response)
|
||||||
|
)
|
||||||
|
})?;
|
||||||
let expires_at = Utc::now().timestamp() + expires_in;
|
let expires_at = Utc::now().timestamp() + expires_in;
|
||||||
|
|
||||||
let tokens = OAuthTokens {
|
let tokens = OAuthTokens {
|
||||||
@@ -363,11 +395,6 @@ async fn run_client_credentials_flow(
|
|||||||
account_id: provider.extract_account_id(&response),
|
account_id: provider.extract_account_id(&response),
|
||||||
};
|
};
|
||||||
save_oauth_tokens(client_name, &tokens)?;
|
save_oauth_tokens(client_name, &tokens)?;
|
||||||
println!(
|
|
||||||
"Successfully authenticated client '{}' with {} via OAuth (client_credentials). Tokens saved.",
|
|
||||||
client_name,
|
|
||||||
provider.provider_name()
|
|
||||||
);
|
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
@@ -417,19 +444,28 @@ async fn run_device_code_flow(provider: &dyn OAuthProvider, client_name: &str) -
|
|||||||
let device_code = device_response["device_code"]
|
let device_code = device_response["device_code"]
|
||||||
.as_str()
|
.as_str()
|
||||||
.ok_or_else(|| {
|
.ok_or_else(|| {
|
||||||
anyhow!("Missing device_code in device authorization response: {device_response}")
|
anyhow!(
|
||||||
|
"Missing device_code in device authorization response (keys: {})",
|
||||||
|
token_response_keys(&device_response)
|
||||||
|
)
|
||||||
})?
|
})?
|
||||||
.to_string();
|
.to_string();
|
||||||
let user_code = device_response["user_code"]
|
let user_code = device_response["user_code"]
|
||||||
.as_str()
|
.as_str()
|
||||||
.ok_or_else(|| {
|
.ok_or_else(|| {
|
||||||
anyhow!("Missing user_code in device authorization response: {device_response}")
|
anyhow!(
|
||||||
|
"Missing user_code in device authorization response (keys: {})",
|
||||||
|
token_response_keys(&device_response)
|
||||||
|
)
|
||||||
})?
|
})?
|
||||||
.to_string();
|
.to_string();
|
||||||
let verification_uri = device_response["verification_uri"]
|
let verification_uri = device_response["verification_uri"]
|
||||||
.as_str()
|
.as_str()
|
||||||
.ok_or_else(|| {
|
.ok_or_else(|| {
|
||||||
anyhow!("Missing verification_uri in device authorization response: {device_response}")
|
anyhow!(
|
||||||
|
"Missing verification_uri in device authorization response (keys: {})",
|
||||||
|
token_response_keys(&device_response)
|
||||||
|
)
|
||||||
})?
|
})?
|
||||||
.to_string();
|
.to_string();
|
||||||
let verification_uri_complete = device_response["verification_uri_complete"]
|
let verification_uri_complete = device_response["verification_uri_complete"]
|
||||||
@@ -551,10 +587,76 @@ fn save_oauth_tokens(client_name: &str, tokens: &OAuthTokens) -> Result<()> {
|
|||||||
fs::create_dir_all(parent)?;
|
fs::create_dir_all(parent)?;
|
||||||
}
|
}
|
||||||
let json = serde_json::to_string_pretty(tokens)?;
|
let json = serde_json::to_string_pretty(tokens)?;
|
||||||
fs::write(path, json)?;
|
// Write-then-rename so a crash mid-write never truncates the live token file.
|
||||||
|
let mut tmp = path.clone().into_os_string();
|
||||||
|
tmp.push(".tmp");
|
||||||
|
let tmp = PathBuf::from(tmp);
|
||||||
|
// Tokens are live credentials: create the file owner-only, not umask-default.
|
||||||
|
let mut options = fs::OpenOptions::new();
|
||||||
|
options.write(true).create(true).truncate(true);
|
||||||
|
#[cfg(unix)]
|
||||||
|
{
|
||||||
|
use std::os::unix::fs::OpenOptionsExt;
|
||||||
|
options.mode(0o600);
|
||||||
|
}
|
||||||
|
options.open(&tmp)?.write_all(json.as_bytes())?;
|
||||||
|
fs::rename(&tmp, &path)?;
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(crate) fn token_response_keys(response: &Value) -> String {
|
||||||
|
match response.as_object() {
|
||||||
|
Some(map) => {
|
||||||
|
let keys: Vec<&str> = map.keys().map(String::as_str).collect();
|
||||||
|
format!("[{}]", keys.join(", "))
|
||||||
|
}
|
||||||
|
None => "<non-object response>".to_string(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_refresh_response(
|
||||||
|
status: StatusCode,
|
||||||
|
response: &Value,
|
||||||
|
previous_refresh_token: Option<&str>,
|
||||||
|
) -> Result<(String, Option<String>, i64)> {
|
||||||
|
if let Some(error) = response["error"].as_str() {
|
||||||
|
let description = response["error_description"]
|
||||||
|
.as_str()
|
||||||
|
.unwrap_or("no description");
|
||||||
|
if matches!(error, "invalid_grant" | "invalid_token") {
|
||||||
|
bail!(
|
||||||
|
"OAuth refresh token was rejected ({error}: {description}). Please re-authenticate."
|
||||||
|
);
|
||||||
|
}
|
||||||
|
bail!("Token refresh failed ({error}: {description})");
|
||||||
|
}
|
||||||
|
if !status.is_success() {
|
||||||
|
bail!("Token refresh failed with HTTP status {status}");
|
||||||
|
}
|
||||||
|
|
||||||
|
let access_token = response["access_token"]
|
||||||
|
.as_str()
|
||||||
|
.ok_or_else(|| {
|
||||||
|
anyhow!(
|
||||||
|
"Missing access_token in refresh response (keys: {})",
|
||||||
|
token_response_keys(response)
|
||||||
|
)
|
||||||
|
})?
|
||||||
|
.to_string();
|
||||||
|
let refresh_token = response["refresh_token"]
|
||||||
|
.as_str()
|
||||||
|
.or(previous_refresh_token)
|
||||||
|
.map(str::to_string);
|
||||||
|
let expires_in = response["expires_in"].as_i64().ok_or_else(|| {
|
||||||
|
anyhow!(
|
||||||
|
"Missing expires_in in refresh response (keys: {})",
|
||||||
|
token_response_keys(response)
|
||||||
|
)
|
||||||
|
})?;
|
||||||
|
|
||||||
|
Ok((access_token, refresh_token, expires_in))
|
||||||
|
}
|
||||||
|
|
||||||
pub async fn refresh_oauth_token(
|
pub async fn refresh_oauth_token(
|
||||||
client: &ReqwestClient,
|
client: &ReqwestClient,
|
||||||
provider: &dyn OAuthProvider,
|
provider: &dyn OAuthProvider,
|
||||||
@@ -577,19 +679,23 @@ pub async fn refresh_oauth_token(
|
|||||||
],
|
],
|
||||||
);
|
);
|
||||||
|
|
||||||
let response: Value = request.send().await?.json().await?;
|
let (status, response) = tokio::time::timeout(TOKEN_ENDPOINT_TIMEOUT, async {
|
||||||
|
let response = request.send().await?;
|
||||||
|
let status = response.status();
|
||||||
|
let body: Value = response.json().await?;
|
||||||
|
Ok::<_, Error>((status, body))
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.map_err(|_| {
|
||||||
|
anyhow!(
|
||||||
|
"Token refresh for '{}' timed out after {}s",
|
||||||
|
client_name,
|
||||||
|
TOKEN_ENDPOINT_TIMEOUT.as_secs()
|
||||||
|
)
|
||||||
|
})??;
|
||||||
|
|
||||||
let access_token = response["access_token"]
|
let (access_token, refresh_token, expires_in) =
|
||||||
.as_str()
|
parse_refresh_response(status, &response, tokens.refresh_token.as_deref())?;
|
||||||
.ok_or_else(|| anyhow!("Missing access_token in refresh response: {response}"))?
|
|
||||||
.to_string();
|
|
||||||
let refresh_token = response["refresh_token"]
|
|
||||||
.as_str()
|
|
||||||
.map(|s| s.to_string())
|
|
||||||
.or_else(|| tokens.refresh_token.clone());
|
|
||||||
let expires_in = response["expires_in"]
|
|
||||||
.as_i64()
|
|
||||||
.ok_or_else(|| anyhow!("Missing expires_in in refresh response: {response}"))?;
|
|
||||||
|
|
||||||
let expires_at = Utc::now().timestamp() + expires_in;
|
let expires_at = Utc::now().timestamp() + expires_in;
|
||||||
|
|
||||||
@@ -609,6 +715,20 @@ pub async fn refresh_oauth_token(
|
|||||||
Ok(new_tokens)
|
Ok(new_tokens)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Per-client lock so concurrent requests perform a single refresh.
|
||||||
|
/// Returns a clone of the Arc so the parking_lot guard is dropped before the
|
||||||
|
/// caller awaits on the tokio mutex.
|
||||||
|
fn refresh_guard(client_name: &str) -> Arc<sync::Mutex<()>> {
|
||||||
|
static GUARDS: OnceLock<parking_lot::Mutex<HashMap<String, Arc<sync::Mutex<()>>>>> =
|
||||||
|
OnceLock::new();
|
||||||
|
GUARDS
|
||||||
|
.get_or_init(Default::default)
|
||||||
|
.lock()
|
||||||
|
.entry(client_name.to_string())
|
||||||
|
.or_default()
|
||||||
|
.clone()
|
||||||
|
}
|
||||||
|
|
||||||
pub async fn prepare_oauth_access_token(
|
pub async fn prepare_oauth_access_token(
|
||||||
client: &ReqwestClient,
|
client: &ReqwestClient,
|
||||||
provider: &dyn OAuthProvider,
|
provider: &dyn OAuthProvider,
|
||||||
@@ -623,16 +743,40 @@ pub async fn prepare_oauth_access_token(
|
|||||||
None => return Ok(false),
|
None => return Ok(false),
|
||||||
};
|
};
|
||||||
|
|
||||||
let tokens = if Utc::now().timestamp() >= tokens.expires_at {
|
let tokens = if Utc::now().timestamp() >= tokens.expires_at
|
||||||
match provider.flow() {
|
|| is_rejected(client_name, &tokens.access_token)
|
||||||
OAuthFlow::Pkce | OAuthFlow::DeviceCode => {
|
{
|
||||||
refresh_oauth_token(client, provider, client_name, &tokens).await?
|
let guard = refresh_guard(client_name);
|
||||||
}
|
let _guard = guard.lock().await;
|
||||||
OAuthFlow::ClientCredentials => {
|
|
||||||
run_client_credentials_flow(provider, client_name).await?;
|
// A concurrent caller may have refreshed while we waited for the
|
||||||
load_oauth_tokens(client_name)
|
// lock; a valid in-memory token means the winner already populated
|
||||||
.ok_or_else(|| anyhow!("Token file missing after client_credentials refresh"))?
|
// the cache.
|
||||||
|
if is_valid_access_token(client_name) {
|
||||||
|
return Ok(true);
|
||||||
|
}
|
||||||
|
|
||||||
|
let tokens = match load_oauth_tokens(client_name) {
|
||||||
|
Some(t) => t,
|
||||||
|
None => return Ok(false),
|
||||||
|
};
|
||||||
|
|
||||||
|
if Utc::now().timestamp() >= tokens.expires_at
|
||||||
|
|| is_rejected(client_name, &tokens.access_token)
|
||||||
|
{
|
||||||
|
match provider.flow() {
|
||||||
|
OAuthFlow::Pkce | OAuthFlow::DeviceCode => {
|
||||||
|
refresh_oauth_token(client, provider, client_name, &tokens).await?
|
||||||
|
}
|
||||||
|
OAuthFlow::ClientCredentials => {
|
||||||
|
run_client_credentials_flow(provider, client_name).await?;
|
||||||
|
load_oauth_tokens(client_name).ok_or_else(|| {
|
||||||
|
anyhow!("Token file missing after client_credentials refresh")
|
||||||
|
})?
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
} else {
|
||||||
|
tokens
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
tokens
|
tokens
|
||||||
@@ -644,6 +788,9 @@ pub async fn prepare_oauth_access_token(
|
|||||||
tokens.expires_at,
|
tokens.expires_at,
|
||||||
tokens.account_id,
|
tokens.account_id,
|
||||||
);
|
);
|
||||||
|
// Clear even when the refresh returned the same token (some IdPs reuse
|
||||||
|
// JWTs within validity); otherwise every request re-hits the token endpoint.
|
||||||
|
clear_rejected(client_name);
|
||||||
|
|
||||||
Ok(true)
|
Ok(true)
|
||||||
}
|
}
|
||||||
@@ -886,11 +1033,55 @@ pub(crate) fn client_config_info(
|
|||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
|
use std::ffi::OsString;
|
||||||
|
use std::path::PathBuf;
|
||||||
use std::str;
|
use std::str;
|
||||||
|
use std::time::UNIX_EPOCH;
|
||||||
|
|
||||||
use super::*;
|
use super::*;
|
||||||
|
use crate::client::access_token::{distrust_access_token, get_access_token};
|
||||||
use crate::client::openai_compatible::OpenAICompatibleConfig;
|
use crate::client::openai_compatible::OpenAICompatibleConfig;
|
||||||
use crate::client::{ModelData, ProviderModels};
|
use crate::client::{ModelData, ProviderModels};
|
||||||
|
use crate::utils::get_env_name;
|
||||||
|
use serial_test::serial;
|
||||||
|
use std::{env, time::SystemTime};
|
||||||
|
|
||||||
|
fn with_temp_cache<F: FnOnce()>(f: F) {
|
||||||
|
struct Restore {
|
||||||
|
key: String,
|
||||||
|
prev: Option<OsString>,
|
||||||
|
root: PathBuf,
|
||||||
|
}
|
||||||
|
impl Drop for Restore {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
unsafe {
|
||||||
|
match self.prev.take() {
|
||||||
|
Some(v) => env::set_var(&self.key, v),
|
||||||
|
None => env::remove_var(&self.key),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let _ = fs::remove_dir_all(&self.root);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let unique = SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap()
|
||||||
|
.as_nanos();
|
||||||
|
let root = env::temp_dir().join(format!("coyote-client-oauth-test-{unique}"));
|
||||||
|
fs::create_dir_all(&root).unwrap();
|
||||||
|
let env_key = get_env_name("cache_dir");
|
||||||
|
let prev = env::var_os(&env_key);
|
||||||
|
unsafe {
|
||||||
|
env::set_var(&env_key, &root);
|
||||||
|
}
|
||||||
|
let _restore = Restore {
|
||||||
|
key: env_key,
|
||||||
|
prev,
|
||||||
|
root,
|
||||||
|
};
|
||||||
|
f();
|
||||||
|
}
|
||||||
|
|
||||||
fn base_config() -> OAuthConfig {
|
fn base_config() -> OAuthConfig {
|
||||||
OAuthConfig {
|
OAuthConfig {
|
||||||
@@ -1468,4 +1659,180 @@ scopes:
|
|||||||
"body missing grant_type param: {body}"
|
"body missing grant_type param: {body}"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[serial]
|
||||||
|
fn save_oauth_tokens_roundtrips_and_leaves_no_tmp_file() {
|
||||||
|
with_temp_cache(|| {
|
||||||
|
let tokens = OAuthTokens {
|
||||||
|
access_token: "at-123".into(),
|
||||||
|
refresh_token: Some("rt-456".into()),
|
||||||
|
expires_at: 1234567890,
|
||||||
|
account_id: Some("acct-789".into()),
|
||||||
|
};
|
||||||
|
|
||||||
|
save_oauth_tokens("atomic-test", &tokens).unwrap();
|
||||||
|
|
||||||
|
let loaded = load_oauth_tokens("atomic-test").unwrap();
|
||||||
|
assert_eq!(loaded.access_token, "at-123");
|
||||||
|
assert_eq!(loaded.refresh_token.as_deref(), Some("rt-456"));
|
||||||
|
assert_eq!(loaded.expires_at, 1234567890);
|
||||||
|
assert_eq!(loaded.account_id.as_deref(), Some("acct-789"));
|
||||||
|
|
||||||
|
let dir = paths::oauth_tokens_dir();
|
||||||
|
let leftover_tmp = fs::read_dir(&dir)
|
||||||
|
.unwrap()
|
||||||
|
.any(|e| e.unwrap().file_name().to_string_lossy().ends_with(".tmp"));
|
||||||
|
assert!(!leftover_tmp, "temp file left behind in {dir:?}");
|
||||||
|
|
||||||
|
#[cfg(unix)]
|
||||||
|
{
|
||||||
|
use std::os::unix::fs::PermissionsExt;
|
||||||
|
let mode = fs::metadata(paths::token_file("atomic-test"))
|
||||||
|
.unwrap()
|
||||||
|
.permissions()
|
||||||
|
.mode();
|
||||||
|
assert_eq!(mode & 0o777, 0o600, "token file mode was {mode:o}");
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[serial]
|
||||||
|
fn prepare_rejected_valid_file_token_attempts_refresh_branch() {
|
||||||
|
with_temp_cache(|| {
|
||||||
|
let client_name = "prepare-rejected-branch-test";
|
||||||
|
let expires_at = Utc::now().timestamp() + 3600;
|
||||||
|
save_oauth_tokens(
|
||||||
|
client_name,
|
||||||
|
&OAuthTokens {
|
||||||
|
access_token: "rejected-at".into(),
|
||||||
|
refresh_token: None,
|
||||||
|
expires_at,
|
||||||
|
account_id: None,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
set_access_token(client_name, "rejected-at".into(), expires_at, None);
|
||||||
|
assert!(distrust_access_token(client_name, "rejected-at"));
|
||||||
|
|
||||||
|
let err = tokio::runtime::Builder::new_current_thread()
|
||||||
|
.enable_all()
|
||||||
|
.build()
|
||||||
|
.unwrap()
|
||||||
|
.block_on(prepare_oauth_access_token(
|
||||||
|
&ReqwestClient::new(),
|
||||||
|
&ResourceStubProvider,
|
||||||
|
client_name,
|
||||||
|
))
|
||||||
|
.unwrap_err()
|
||||||
|
.to_string();
|
||||||
|
|
||||||
|
// The timestamp-valid but rejected file token must not be trusted;
|
||||||
|
// the refresh branch is taken and bails on the missing refresh token.
|
||||||
|
assert!(err.contains("No refresh token"), "unexpected error: {err}");
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[serial]
|
||||||
|
fn prepare_trusts_differing_unmarked_valid_file_token() {
|
||||||
|
with_temp_cache(|| {
|
||||||
|
let client_name = "prepare-differing-token-test";
|
||||||
|
let expires_at = Utc::now().timestamp() + 3600;
|
||||||
|
set_access_token(client_name, "rejected-at".into(), expires_at, None);
|
||||||
|
assert!(distrust_access_token(client_name, "rejected-at"));
|
||||||
|
save_oauth_tokens(
|
||||||
|
client_name,
|
||||||
|
&OAuthTokens {
|
||||||
|
access_token: "fresh-at".into(),
|
||||||
|
refresh_token: None,
|
||||||
|
expires_at,
|
||||||
|
account_id: None,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let ready = tokio::runtime::Builder::new_current_thread()
|
||||||
|
.enable_all()
|
||||||
|
.build()
|
||||||
|
.unwrap()
|
||||||
|
.block_on(prepare_oauth_access_token(
|
||||||
|
&ReqwestClient::new(),
|
||||||
|
&ResourceStubProvider,
|
||||||
|
client_name,
|
||||||
|
))
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert!(ready);
|
||||||
|
assert_eq!(get_access_token(client_name).unwrap(), "fresh-at");
|
||||||
|
assert!(
|
||||||
|
!is_rejected(client_name, "rejected-at"),
|
||||||
|
"marker not cleared after successful prepare"
|
||||||
|
);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_refresh_response_invalid_grant_redacts_and_prompts_reauth() {
|
||||||
|
let response = serde_json::json!({
|
||||||
|
"error": "invalid_grant",
|
||||||
|
"error_description": "refresh token revoked",
|
||||||
|
"refresh_token": "planted-secret-token",
|
||||||
|
});
|
||||||
|
|
||||||
|
let err = parse_refresh_response(StatusCode::BAD_REQUEST, &response, Some("old-rt"))
|
||||||
|
.unwrap_err()
|
||||||
|
.to_string();
|
||||||
|
|
||||||
|
assert!(err.contains("re-authenticate"), "unexpected error: {err}");
|
||||||
|
assert!(
|
||||||
|
!err.contains("planted-secret-token"),
|
||||||
|
"error leaked token material: {err}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn token_response_keys_lists_keys_without_values() {
|
||||||
|
let response = serde_json::json!({
|
||||||
|
"access_token": "secret-at",
|
||||||
|
"token_type": "SecretBearer",
|
||||||
|
});
|
||||||
|
|
||||||
|
let keys = token_response_keys(&response);
|
||||||
|
|
||||||
|
assert!(keys.contains("access_token"), "missing key name: {keys}");
|
||||||
|
assert!(keys.contains("token_type"), "missing key name: {keys}");
|
||||||
|
assert!(!keys.contains("secret-at"), "leaked value: {keys}");
|
||||||
|
assert!(!keys.contains("SecretBearer"), "leaked value: {keys}");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_refresh_response_rotates_refresh_token_when_present() {
|
||||||
|
let response = serde_json::json!({
|
||||||
|
"access_token": "new-at",
|
||||||
|
"refresh_token": "new-rt",
|
||||||
|
"expires_in": 3600,
|
||||||
|
});
|
||||||
|
|
||||||
|
let (access_token, refresh_token, expires_in) =
|
||||||
|
parse_refresh_response(StatusCode::OK, &response, Some("old-rt")).unwrap();
|
||||||
|
|
||||||
|
assert_eq!(access_token, "new-at");
|
||||||
|
assert_eq!(refresh_token.as_deref(), Some("new-rt"));
|
||||||
|
assert_eq!(expires_in, 3600);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_refresh_response_keeps_old_refresh_token_when_absent() {
|
||||||
|
let response = serde_json::json!({
|
||||||
|
"access_token": "new-at",
|
||||||
|
"expires_in": 3600,
|
||||||
|
});
|
||||||
|
|
||||||
|
let (_, refresh_token, _) =
|
||||||
|
parse_refresh_response(StatusCode::OK, &response, Some("old-rt")).unwrap();
|
||||||
|
|
||||||
|
assert_eq!(refresh_token.as_deref(), Some("old-rt"));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+54
-5
@@ -1,4 +1,4 @@
|
|||||||
use super::{ThinkingBlock, ToolCall, catch_error};
|
use super::{ApiStatusError, ThinkingBlock, ToolCall, catch_error};
|
||||||
use crate::utils::AbortSignal;
|
use crate::utils::AbortSignal;
|
||||||
|
|
||||||
use anyhow::{Context, Result, anyhow, bail};
|
use anyhow::{Context, Result, anyhow, bail};
|
||||||
@@ -176,6 +176,14 @@ impl SseHandler {
|
|||||||
self.thinking.push(block);
|
self.thinking.push(block);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Whether any output (text, tool calls, or thinking blocks) has been
|
||||||
|
/// accumulated. `Client::chat_completions_streaming` gates its 401 retry
|
||||||
|
/// on this: content already streamed to the user would be rendered a
|
||||||
|
/// second time by a retry, so partial responses are never retried.
|
||||||
|
pub fn has_received_content(&self) -> bool {
|
||||||
|
!self.buffer.is_empty() || !self.tool_calls.is_empty() || !self.thinking.is_empty()
|
||||||
|
}
|
||||||
|
|
||||||
pub fn abort(&self) -> AbortSignal {
|
pub fn abort(&self) -> AbortSignal {
|
||||||
self.abort_signal.clone()
|
self.abort_signal.clone()
|
||||||
}
|
}
|
||||||
@@ -224,10 +232,14 @@ where
|
|||||||
let data: Value = match text.parse() {
|
let data: Value = match text.parse() {
|
||||||
Ok(data) => data,
|
Ok(data) => data,
|
||||||
Err(_) => {
|
Err(_) => {
|
||||||
bail!(
|
return Err(ApiStatusError {
|
||||||
"Invalid response data: {text} (status: {})",
|
status: status.as_u16(),
|
||||||
status.as_u16()
|
message: format!(
|
||||||
);
|
"Invalid response data: {text} (status: {})",
|
||||||
|
status.as_u16()
|
||||||
|
),
|
||||||
|
}
|
||||||
|
.into());
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
catch_error(&data, status.as_u16())?;
|
catch_error(&data, status.as_u16())?;
|
||||||
@@ -418,6 +430,43 @@ mod tests {
|
|||||||
assert!(error_message.contains("test_function_loop"));
|
assert!(error_message.contains("test_function_loop"));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn new_handler() -> (SseHandler, tokio::sync::mpsc::UnboundedReceiver<SseEvent>) {
|
||||||
|
let (sender, receiver) = tokio::sync::mpsc::unbounded_channel();
|
||||||
|
let abort_signal = crate::utils::create_abort_signal();
|
||||||
|
(SseHandler::new(sender, abort_signal), receiver)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_has_received_content_text() {
|
||||||
|
let (mut handler, _rx) = new_handler();
|
||||||
|
assert!(!handler.has_received_content());
|
||||||
|
|
||||||
|
handler.text("hello").unwrap();
|
||||||
|
assert!(handler.has_received_content());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_has_received_content_tool_call() {
|
||||||
|
let (mut handler, _rx) = new_handler();
|
||||||
|
assert!(!handler.has_received_content());
|
||||||
|
|
||||||
|
let call = ToolCall::new("test_function".to_string(), json!({"param": 1}), None);
|
||||||
|
handler.tool_call(call).unwrap();
|
||||||
|
assert!(handler.has_received_content());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_has_received_content_thinking() {
|
||||||
|
let (mut handler, _rx) = new_handler();
|
||||||
|
assert!(!handler.has_received_content());
|
||||||
|
|
||||||
|
handler.thinking_block(ThinkingBlock::Thinking {
|
||||||
|
thinking: "hmm".to_string(),
|
||||||
|
signature: "sig".to_string(),
|
||||||
|
});
|
||||||
|
assert!(handler.has_received_content());
|
||||||
|
}
|
||||||
|
|
||||||
fn split_chunks(text: &str) -> Vec<Vec<u8>> {
|
fn split_chunks(text: &str) -> Vec<Vec<u8>> {
|
||||||
let len = text.len();
|
let len = text.len();
|
||||||
let cut1 = random_range(1..len - 1);
|
let cut1 = random_range(1..len - 1);
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
use crate::client::{ClientConfig, Model, ModelType, list_models};
|
use crate::client::{ClientConfig, Model, ModelType, list_models};
|
||||||
use crate::render::{MarkdownRender, RenderOptions};
|
use crate::render::{MarkdownRender, RenderOptions};
|
||||||
use crate::utils::{IS_STDOUT_TERMINAL, NO_COLOR, decode_bin, get_env_name};
|
use crate::utils::{IS_STDOUT_TERMINAL, NO_COLOR, decode_bin, drain_stale_tty_input, get_env_name};
|
||||||
|
|
||||||
use super::paths;
|
use super::paths;
|
||||||
use anyhow::{Context, Result, anyhow, bail};
|
use anyhow::{Context, Result, anyhow, bail};
|
||||||
@@ -616,14 +616,18 @@ impl AppConfig {
|
|||||||
if self.highlight && self.theme.is_none() {
|
if self.highlight && self.theme.is_none() {
|
||||||
if let Some(v) = super::read_env_value::<String>(&get_env_name("theme")) {
|
if let Some(v) = super::read_env_value::<String>(&get_env_name("theme")) {
|
||||||
self.theme = v;
|
self.theme = v;
|
||||||
} else if *IS_STDOUT_TERMINAL
|
} else if *IS_STDOUT_TERMINAL {
|
||||||
&& let Ok(color_scheme) = color_scheme(QueryOptions::default())
|
if let Ok(color_scheme) = color_scheme(QueryOptions::default()) {
|
||||||
{
|
let theme = match color_scheme {
|
||||||
let theme = match color_scheme {
|
ColorScheme::Dark => "dark",
|
||||||
ColorScheme::Dark => "dark",
|
ColorScheme::Light => "light",
|
||||||
ColorScheme::Light => "light",
|
};
|
||||||
};
|
self.theme = Some(theme.into());
|
||||||
self.theme = Some(theme.into());
|
}
|
||||||
|
// The OSC/DA1 reply can arrive after colorsaurus stops reading
|
||||||
|
// (observed under zellij-in-kitty). Drain any late reply bytes so
|
||||||
|
// they are neither echoed nor read as line-editor input.
|
||||||
|
drain_stale_tty_input();
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if let Some(v) = super::read_env_value::<String>(&get_env_name("left_prompt")) {
|
if let Some(v) = super::read_env_value::<String>(&get_env_name("left_prompt")) {
|
||||||
|
|||||||
+13
-19
@@ -1,6 +1,6 @@
|
|||||||
use crate::mcp::{
|
use crate::mcp::{
|
||||||
ConnectedServer, JsonField, McpServer, McpTransportType, is_auth_required_error, oauth,
|
ConnectedServer, JsonField, McpAuthRequired, McpServer, McpTransportType,
|
||||||
spawn_mcp_server,
|
is_auth_required_error, resolve_http_auth, spawn_mcp_server,
|
||||||
};
|
};
|
||||||
|
|
||||||
use anyhow::Result;
|
use anyhow::Result;
|
||||||
@@ -102,23 +102,17 @@ impl McpFactory {
|
|||||||
return Ok(existing);
|
return Ok(existing);
|
||||||
}
|
}
|
||||||
|
|
||||||
let bearer_token = if spec.is_remote() {
|
let (auth, auth_reason) = resolve_http_auth(name, spec).await;
|
||||||
oauth::load_valid_mcp_token(name)
|
let handle = spawn_mcp_server(spec, log_path, auth).await.map_err(|e| {
|
||||||
} else {
|
if is_auth_required_error(&e) {
|
||||||
None
|
e.context(McpAuthRequired {
|
||||||
};
|
server: name.to_string(),
|
||||||
let handle = spawn_mcp_server(spec, log_path, bearer_token)
|
reason: auth_reason,
|
||||||
.await
|
})
|
||||||
.map_err(|e| {
|
} else {
|
||||||
if is_auth_required_error(&e) {
|
e
|
||||||
e.context(format!(
|
}
|
||||||
"MCP server '{name}' requires OAuth authentication. \
|
})?;
|
||||||
Run `coyote --auth-mcp {name}` or `.mcp auth {name}` in the REPL to authenticate."
|
|
||||||
))
|
|
||||||
} else {
|
|
||||||
e
|
|
||||||
}
|
|
||||||
})?;
|
|
||||||
self.insert_active(key, &handle);
|
self.insert_active(key, &handle);
|
||||||
Ok(handle)
|
Ok(handle)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -148,6 +148,16 @@ pub fn sbx_kit_hash_file() -> PathBuf {
|
|||||||
sbx_kit_dir().join(SBX_KIT_HASH_FILE)
|
sbx_kit_dir().join(SBX_KIT_HASH_FILE)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn sandbox_mixin_hashes_dir() -> PathBuf {
|
||||||
|
cache_dir().join("sandbox-mixin-hashes")
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn sandbox_mixin_hash_file(sandbox_name: &str) -> PathBuf {
|
||||||
|
// Sandbox names are sanitized by the caller, but never trust a path
|
||||||
|
// component: a stray separator must not escape the hash directory.
|
||||||
|
sandbox_mixin_hashes_dir().join(format!("{}.hash", sandbox_name.replace('/', "_")))
|
||||||
|
}
|
||||||
|
|
||||||
pub fn sbx_mixin_kits_dir() -> PathBuf {
|
pub fn sbx_mixin_kits_dir() -> PathBuf {
|
||||||
cache_dir().join(SBX_MIXIN_KITS_DIR_NAME)
|
cache_dir().join(SBX_MIXIN_KITS_DIR_NAME)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -22,7 +22,7 @@ use crate::function::{
|
|||||||
};
|
};
|
||||||
use crate::mcp::{
|
use crate::mcp::{
|
||||||
MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, MCP_INVOKE_META_FUNCTION_NAME_PREFIX,
|
MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, MCP_INVOKE_META_FUNCTION_NAME_PREFIX,
|
||||||
MCP_SEARCH_META_FUNCTION_NAME_PREFIX, is_auth_required_error,
|
MCP_SEARCH_META_FUNCTION_NAME_PREFIX, McpAuthReason, McpAuthRequired, is_auth_required_error,
|
||||||
};
|
};
|
||||||
use crate::rag::Rag;
|
use crate::rag::Rag;
|
||||||
use crate::supervisor::Supervisor;
|
use crate::supervisor::Supervisor;
|
||||||
@@ -3427,7 +3427,11 @@ impl RequestContext {
|
|||||||
{
|
{
|
||||||
Ok(handle) => handles.push((id.clone(), handle)),
|
Ok(handle) => handles.push((id.clone(), handle)),
|
||||||
Err(e) if is_auth_required_error(&e) => {
|
Err(e) if is_auth_required_error(&e) => {
|
||||||
auth_required.push(id.clone())
|
let reason = e
|
||||||
|
.downcast_ref::<McpAuthRequired>()
|
||||||
|
.map(|a| a.reason)
|
||||||
|
.unwrap_or(McpAuthReason::NotAuthenticated);
|
||||||
|
auth_required.push((id.clone(), reason));
|
||||||
}
|
}
|
||||||
Err(e) => return Err(e),
|
Err(e) => return Err(e),
|
||||||
}
|
}
|
||||||
@@ -3444,11 +3448,8 @@ impl RequestContext {
|
|||||||
for (id, handle) in handles {
|
for (id, handle) in handles {
|
||||||
mcp_runtime.insert(id, handle);
|
mcp_runtime.insert(id, handle);
|
||||||
}
|
}
|
||||||
for id in auth_required {
|
for (id, reason) in auth_required {
|
||||||
eprintln!(
|
eprintln!("Warning: {}", McpAuthRequired { server: id, reason });
|
||||||
"Warning: MCP server '{id}' requires OAuth authentication and was not started. \
|
|
||||||
Run `.mcp auth {id}` (or `coyote --auth-mcp {id}`) to authenticate and attach it."
|
|
||||||
);
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+445
-97
@@ -29,16 +29,17 @@ use rust_embed::Embed;
|
|||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use serde_json::{Value, json};
|
use serde_json::{Value, json};
|
||||||
use skill::SKILL_FUNCTION_PREFIX;
|
use skill::SKILL_FUNCTION_PREFIX;
|
||||||
use std::collections::VecDeque;
|
|
||||||
use std::ffi::OsStr;
|
use std::ffi::OsStr;
|
||||||
use std::fs::File;
|
use std::fs::File;
|
||||||
use std::io::{Read, Write};
|
use std::io::{Read, Write};
|
||||||
use std::sync::atomic::Ordering;
|
use std::sync::atomic::{AtomicU64, Ordering};
|
||||||
|
use std::{collections::VecDeque, thread};
|
||||||
use std::{
|
use std::{
|
||||||
collections::{HashMap, HashSet},
|
collections::{HashMap, HashSet},
|
||||||
env, fs, io,
|
env, fs, io,
|
||||||
path::{Path, PathBuf},
|
path::{Path, PathBuf},
|
||||||
process::{Command, Stdio},
|
process::{Command, Stdio},
|
||||||
|
time::{Duration, Instant},
|
||||||
};
|
};
|
||||||
use strum_macros::AsRefStr;
|
use strum_macros::AsRefStr;
|
||||||
use supervisor::SUPERVISOR_FUNCTION_PREFIX;
|
use supervisor::SUPERVISOR_FUNCTION_PREFIX;
|
||||||
@@ -137,6 +138,114 @@ fn extract_shebang_runtime(path: &Path) -> Option<String> {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(crate) fn write_file_atomic(
|
||||||
|
path: &Path,
|
||||||
|
content: &str,
|
||||||
|
#[cfg_attr(not(unix), expect(unused))] mode: Option<u32>,
|
||||||
|
) -> Result<()> {
|
||||||
|
#[cfg(unix)]
|
||||||
|
use std::os::unix::fs::PermissionsExt;
|
||||||
|
|
||||||
|
if fs::read_to_string(path).is_ok_and(|existing| existing == content) {
|
||||||
|
#[cfg(unix)]
|
||||||
|
if let Some(mode) = mode {
|
||||||
|
fs::set_permissions(path, fs::Permissions::from_mode(mode))?;
|
||||||
|
}
|
||||||
|
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
|
let file_name = path
|
||||||
|
.file_name()
|
||||||
|
.and_then(OsStr::to_str)
|
||||||
|
.ok_or_else(|| anyhow!("Unable to extract file name from path: {}", path.display()))?;
|
||||||
|
static TMP_COUNTER: AtomicU64 = AtomicU64::new(0);
|
||||||
|
let tmp = path.with_file_name(format!(
|
||||||
|
".{file_name}.tmp.{}.{}",
|
||||||
|
std::process::id(),
|
||||||
|
TMP_COUNTER.fetch_add(1, Ordering::Relaxed)
|
||||||
|
));
|
||||||
|
fs::write(&tmp, content)?;
|
||||||
|
|
||||||
|
#[cfg(unix)]
|
||||||
|
if let Some(mode) = mode {
|
||||||
|
fs::set_permissions(&tmp, fs::Permissions::from_mode(mode))?;
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Err(err) = fs::rename(&tmp, path) {
|
||||||
|
let _ = fs::remove_file(&tmp);
|
||||||
|
return Err(err.into());
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn tool_source_stems() -> Result<HashSet<String>> {
|
||||||
|
let mut stems = HashSet::new();
|
||||||
|
let tools_dir = paths::global_tools_dir();
|
||||||
|
if !tools_dir.exists() {
|
||||||
|
return Ok(stems);
|
||||||
|
}
|
||||||
|
|
||||||
|
for entry in fs::read_dir(&tools_dir)? {
|
||||||
|
let path = entry?.path();
|
||||||
|
if path.is_file()
|
||||||
|
&& let Some(stem) = path.file_stem().and_then(OsStr::to_str)
|
||||||
|
{
|
||||||
|
stems.insert(stem.to_string());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(stems)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bin_entry_stem(file_name: &str) -> &str {
|
||||||
|
let name = file_name.strip_prefix("run-").unwrap_or(file_name);
|
||||||
|
Path::new(name)
|
||||||
|
.file_stem()
|
||||||
|
.and_then(OsStr::to_str)
|
||||||
|
.unwrap_or(name)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn prune_stale_bin_entries(
|
||||||
|
bin_dir: &Path,
|
||||||
|
valid_stems: &HashSet<String>,
|
||||||
|
extra_valid_stem: Option<&str>,
|
||||||
|
) -> Result<()> {
|
||||||
|
if !bin_dir.exists() {
|
||||||
|
fs::create_dir_all(bin_dir)?;
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
|
for entry in fs::read_dir(bin_dir)? {
|
||||||
|
let entry = entry?;
|
||||||
|
let path = entry.path();
|
||||||
|
if path.is_dir() {
|
||||||
|
debug!(
|
||||||
|
"Removing unexpected directory in bin dir: {}",
|
||||||
|
path.display()
|
||||||
|
);
|
||||||
|
fs::remove_dir_all(&path)?;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
let file_name = entry.file_name();
|
||||||
|
let Some(file_name) = file_name.to_str() else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
let stem = bin_entry_stem(file_name);
|
||||||
|
|
||||||
|
if valid_stems.contains(stem) || extra_valid_stem == Some(stem) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
debug!("Removing stale bin entry: {}", path.display());
|
||||||
|
fs::remove_file(&path)?;
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
pub async fn eval_tool_calls(
|
pub async fn eval_tool_calls(
|
||||||
ctx: &mut RequestContext,
|
ctx: &mut RequestContext,
|
||||||
mut calls: Vec<ToolCall>,
|
mut calls: Vec<ToolCall>,
|
||||||
@@ -185,42 +294,29 @@ pub async fn eval_tool_calls(
|
|||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
for (idx, call, result) in future::join_all(futs).await {
|
for (idx, call, result) in future::join_all(futs).await {
|
||||||
indexed_results.push((idx, ToolResult::new(call, normalize_tool_result(result?))));
|
let value = match result {
|
||||||
|
Ok(v) => normalize_tool_result(v),
|
||||||
|
Err(e) => json!({"tool_call_error": format!("{e}")}),
|
||||||
|
};
|
||||||
|
indexed_results.push((idx, ToolResult::new(call, value)));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
for (idx, call) in sequential_calls {
|
for (idx, call) in sequential_calls {
|
||||||
let result = call.eval(ctx).await?;
|
let value = match call.eval(ctx).await {
|
||||||
indexed_results.push((idx, ToolResult::new(call, normalize_tool_result(result))));
|
Ok(v) => normalize_tool_result(v),
|
||||||
|
Err(e) => json!({
|
||||||
|
"tool_call_error": format!(
|
||||||
|
"{e}. This tool is not available or the call failed; use only tools listed in your catalog."
|
||||||
|
)
|
||||||
|
}),
|
||||||
|
};
|
||||||
|
indexed_results.push((idx, ToolResult::new(call, value)));
|
||||||
}
|
}
|
||||||
|
|
||||||
indexed_results.sort_unstable_by_key(|(idx, _)| *idx);
|
indexed_results.sort_unstable_by_key(|(idx, _)| *idx);
|
||||||
output = indexed_results.into_iter().map(|(_, r)| r).collect();
|
output = indexed_results.into_iter().map(|(_, r)| r).collect();
|
||||||
|
|
||||||
if !output.is_empty() {
|
|
||||||
let (has_escalations, summary) = if ctx.current_depth == 0
|
|
||||||
&& let Some(queue) = ctx.root_escalation_queue()
|
|
||||||
&& queue.has_pending()
|
|
||||||
{
|
|
||||||
(true, queue.pending_summary())
|
|
||||||
} else {
|
|
||||||
(false, vec![])
|
|
||||||
};
|
|
||||||
|
|
||||||
if has_escalations {
|
|
||||||
let notification = json!({
|
|
||||||
"pending_escalations": summary,
|
|
||||||
"instruction": "Child agents are BLOCKED waiting for your reply. Call agent__reply_escalation for each pending escalation to unblock them."
|
|
||||||
});
|
|
||||||
let synthetic_call = ToolCall::new(
|
|
||||||
"__escalation_notification".to_string(),
|
|
||||||
json!({}),
|
|
||||||
Some("escalation_check".to_string()),
|
|
||||||
);
|
|
||||||
output.push(ToolResult::new(synthetic_call, notification));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
{
|
{
|
||||||
let max_chars = ctx
|
let max_chars = ctx
|
||||||
.agent
|
.agent
|
||||||
@@ -235,6 +331,14 @@ pub async fn eval_tool_calls(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if ctx.current_depth == 0
|
||||||
|
&& let Some(queue) = ctx.root_escalation_queue()
|
||||||
|
&& queue.has_pending()
|
||||||
|
&& let Some(last) = output.last_mut()
|
||||||
|
{
|
||||||
|
inject_escalation_notification(last, queue.pending_summary());
|
||||||
|
}
|
||||||
|
|
||||||
Ok(output)
|
Ok(output)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -251,6 +355,24 @@ fn normalize_tool_result(result: Value) -> Value {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn inject_escalation_notification(last: &mut ToolResult, summary: Vec<Value>) {
|
||||||
|
let instruction = "Child agents are BLOCKED waiting for your reply. \
|
||||||
|
Call agent__reply_escalation for each pending escalation to unblock them.";
|
||||||
|
match &mut last.output {
|
||||||
|
Value::Object(map) => {
|
||||||
|
map.insert("pending_escalations".into(), json!(summary));
|
||||||
|
map.insert("escalation_instruction".into(), json!(instruction));
|
||||||
|
}
|
||||||
|
other => {
|
||||||
|
*other = json!({
|
||||||
|
"output": other.take(),
|
||||||
|
"pending_escalations": summary,
|
||||||
|
"escalation_instruction": instruction,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||||
pub struct ToolResult {
|
pub struct ToolResult {
|
||||||
pub call: ToolCall,
|
pub call: ToolCall,
|
||||||
@@ -307,7 +429,6 @@ impl Functions {
|
|||||||
})?;
|
})?;
|
||||||
let content = unsafe { std::str::from_utf8_unchecked(&embedded_file.data) };
|
let content = unsafe { std::str::from_utf8_unchecked(&embedded_file.data) };
|
||||||
let file_path = paths::functions_dir().join(file.as_ref());
|
let file_path = paths::functions_dir().join(file.as_ref());
|
||||||
#[cfg_attr(not(unix), expect(unused))]
|
|
||||||
let is_script = file_path
|
let is_script = file_path
|
||||||
.extension()
|
.extension()
|
||||||
.and_then(OsStr::to_str)
|
.and_then(OsStr::to_str)
|
||||||
@@ -324,14 +445,7 @@ impl Functions {
|
|||||||
|
|
||||||
ensure_parent_exists(&file_path)?;
|
ensure_parent_exists(&file_path)?;
|
||||||
info!("Creating function file: {}", file_path.display());
|
info!("Creating function file: {}", file_path.display());
|
||||||
let mut function_file = File::create(&file_path)?;
|
write_file_atomic(&file_path, content, is_script.then_some(0o755))?;
|
||||||
function_file.write_all(content.as_bytes())?;
|
|
||||||
|
|
||||||
#[cfg(unix)]
|
|
||||||
if is_script {
|
|
||||||
use std::os::unix::fs::PermissionsExt;
|
|
||||||
fs::set_permissions(&file_path, fs::Permissions::from_mode(0o755))?;
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
@@ -382,7 +496,7 @@ impl Functions {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub fn init(visible_tools: &[String]) -> Result<Self> {
|
pub fn init(visible_tools: &[String]) -> Result<Self> {
|
||||||
Self::clear_global_functions_bin_dir()?;
|
Self::remove_stale_global_function_binaries()?;
|
||||||
|
|
||||||
let declarations = Self {
|
let declarations = Self {
|
||||||
declarations: Self::build_global_tool_declarations(visible_tools)?,
|
declarations: Self::build_global_tool_declarations(visible_tools)?,
|
||||||
@@ -398,7 +512,7 @@ impl Functions {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub fn init_agent(name: &str, global_tools: &[String]) -> Result<Self> {
|
pub fn init_agent(name: &str, global_tools: &[String]) -> Result<Self> {
|
||||||
Self::clear_agent_bin_dir(name)?;
|
Self::remove_stale_agent_bin_entries(name)?;
|
||||||
|
|
||||||
let global_tools_declarations = if !global_tools.is_empty() {
|
let global_tools_declarations = if !global_tools.is_empty() {
|
||||||
info!("Loading global tools for agent: {name}: {global_tools:?}");
|
info!("Loading global tools for agent: {name}: {global_tools:?}");
|
||||||
@@ -712,38 +826,23 @@ impl Functions {
|
|||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
fn clear_agent_bin_dir(name: &str) -> Result<()> {
|
fn remove_stale_agent_bin_entries(name: &str) -> Result<()> {
|
||||||
let agent_bin_directory = paths::agent_bin_dir(name);
|
let agent_bin_directory = paths::agent_bin_dir(name);
|
||||||
if !agent_bin_directory.exists() {
|
|
||||||
debug!(
|
|
||||||
"Creating agent bin directory: {}",
|
|
||||||
agent_bin_directory.display()
|
|
||||||
);
|
|
||||||
fs::create_dir_all(&agent_bin_directory)?;
|
|
||||||
} else {
|
|
||||||
debug!(
|
|
||||||
"Clearing existing agent bin directory: {}",
|
|
||||||
agent_bin_directory.display()
|
|
||||||
);
|
|
||||||
clear_dir(&agent_bin_directory)?;
|
|
||||||
}
|
|
||||||
|
|
||||||
Ok(())
|
debug!(
|
||||||
|
"Pruning stale entries in agent bin directory: {}",
|
||||||
|
agent_bin_directory.display()
|
||||||
|
);
|
||||||
|
|
||||||
|
prune_stale_bin_entries(&agent_bin_directory, &tool_source_stems()?, Some(name))
|
||||||
}
|
}
|
||||||
|
|
||||||
fn clear_global_functions_bin_dir() -> Result<()> {
|
fn remove_stale_global_function_binaries() -> Result<()> {
|
||||||
let bin_dir = paths::functions_bin_dir();
|
let bin_dir = paths::functions_bin_dir();
|
||||||
if !bin_dir.exists() {
|
|
||||||
fs::create_dir_all(&bin_dir)?;
|
|
||||||
}
|
|
||||||
|
|
||||||
info!(
|
info!("Pruning stale function binaries in {}", bin_dir.display());
|
||||||
"Clearing existing function binaries in {}",
|
|
||||||
bin_dir.display()
|
|
||||||
);
|
|
||||||
clear_dir(&bin_dir)?;
|
|
||||||
|
|
||||||
Ok(())
|
prune_stale_bin_entries(&bin_dir, &tool_source_stems()?, None)
|
||||||
}
|
}
|
||||||
|
|
||||||
fn build_agent_tool_binaries(name: &str) -> Result<()> {
|
fn build_agent_tool_binaries(name: &str) -> Result<()> {
|
||||||
@@ -844,11 +943,7 @@ impl Functions {
|
|||||||
"{prompt_utils_file}",
|
"{prompt_utils_file}",
|
||||||
&to_script_path(&paths::bash_prompt_utils_file().to_string_lossy()),
|
&to_script_path(&paths::bash_prompt_utils_file().to_string_lossy()),
|
||||||
);
|
);
|
||||||
if binary_script_file.exists() {
|
write_file_atomic(&binary_script_file, &content, None)?;
|
||||||
fs::remove_file(&binary_script_file)?;
|
|
||||||
}
|
|
||||||
let mut script_file = File::create(&binary_script_file)?;
|
|
||||||
script_file.write_all(content.as_bytes())?;
|
|
||||||
|
|
||||||
info!(
|
info!(
|
||||||
"Building binary for function: {} ({})",
|
"Building binary for function: {} ({})",
|
||||||
@@ -910,8 +1005,7 @@ impl Functions {
|
|||||||
{run} "{wrapper_binary}" %*"#,
|
{run} "{wrapper_binary}" %*"#,
|
||||||
);
|
);
|
||||||
|
|
||||||
let mut file = File::create(&binary_file)?;
|
write_file_atomic(&binary_file, &content, None)?;
|
||||||
file.write_all(content.as_bytes())?;
|
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
@@ -923,8 +1017,6 @@ impl Functions {
|
|||||||
binary_type: BinaryType,
|
binary_type: BinaryType,
|
||||||
custom_runtime: Option<&str>,
|
custom_runtime: Option<&str>,
|
||||||
) -> Result<()> {
|
) -> Result<()> {
|
||||||
use std::os::unix::prelude::PermissionsExt;
|
|
||||||
|
|
||||||
let binary_file = match binary_type {
|
let binary_file = match binary_type {
|
||||||
BinaryType::Tool(None) => paths::functions_bin_dir().join(binary_name),
|
BinaryType::Tool(None) => paths::functions_bin_dir().join(binary_name),
|
||||||
BinaryType::Tool(Some(agent_name)) => {
|
BinaryType::Tool(Some(agent_name)) => {
|
||||||
@@ -993,31 +1085,16 @@ impl Functions {
|
|||||||
.parent()
|
.parent()
|
||||||
.expect("Failed to get parent directory of binary file");
|
.expect("Failed to get parent directory of binary file");
|
||||||
let script_file = bin_dir.join(format!("run-{binary_name}.ts"));
|
let script_file = bin_dir.join(format!("run-{binary_name}.ts"));
|
||||||
if script_file.exists() {
|
write_file_atomic(&script_file, &content, Some(0o755))?;
|
||||||
fs::remove_file(&script_file)?;
|
|
||||||
}
|
|
||||||
let mut sf = File::create(&script_file)?;
|
|
||||||
sf.write_all(content.as_bytes())?;
|
|
||||||
fs::set_permissions(&script_file, fs::Permissions::from_mode(0o755))?;
|
|
||||||
|
|
||||||
let ts_runtime = custom_runtime.unwrap_or("tsx");
|
let ts_runtime = custom_runtime.unwrap_or("tsx");
|
||||||
let wrapper = format!(
|
let wrapper = format!(
|
||||||
"#!/bin/sh\nexec {ts_runtime} \"{}\" \"$@\"\n",
|
"#!/bin/sh\nexec {ts_runtime} \"{}\" \"$@\"\n",
|
||||||
script_file.display()
|
script_file.display()
|
||||||
);
|
);
|
||||||
if binary_file.exists() {
|
write_file_atomic(&binary_file, &wrapper, Some(0o755))?;
|
||||||
fs::remove_file(&binary_file)?;
|
|
||||||
}
|
|
||||||
let mut wf = File::create(&binary_file)?;
|
|
||||||
wf.write_all(wrapper.as_bytes())?;
|
|
||||||
fs::set_permissions(&binary_file, fs::Permissions::from_mode(0o755))?;
|
|
||||||
} else {
|
} else {
|
||||||
if binary_file.exists() {
|
write_file_atomic(&binary_file, &content, Some(0o755))?;
|
||||||
fs::remove_file(&binary_file)?;
|
|
||||||
}
|
|
||||||
let mut file = File::create(&binary_file)?;
|
|
||||||
file.write_all(content.as_bytes())?;
|
|
||||||
fs::set_permissions(&binary_file, fs::Permissions::from_mode(0o755))?;
|
|
||||||
}
|
}
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
@@ -1456,6 +1533,7 @@ pub fn run_llm_function(
|
|||||||
let mut child = Command::new(&cmd_name)
|
let mut child = Command::new(&cmd_name)
|
||||||
.args(&cmd_args)
|
.args(&cmd_args)
|
||||||
.envs(envs)
|
.envs(envs)
|
||||||
|
.stdin(Stdio::null())
|
||||||
.stdout(Stdio::piped())
|
.stdout(Stdio::piped())
|
||||||
.stderr(Stdio::piped())
|
.stderr(Stdio::piped())
|
||||||
.spawn()
|
.spawn()
|
||||||
@@ -1464,7 +1542,7 @@ pub fn run_llm_function(
|
|||||||
let stdout = child.stdout.take().expect("Failed to capture stdout");
|
let stdout = child.stdout.take().expect("Failed to capture stdout");
|
||||||
let stderr = child.stderr.take().expect("Failed to capture stderr");
|
let stderr = child.stderr.take().expect("Failed to capture stderr");
|
||||||
|
|
||||||
let stdout_thread = std::thread::spawn(move || {
|
let stdout_thread = thread::spawn(move || {
|
||||||
let mut buffer = [0; 1024];
|
let mut buffer = [0; 1024];
|
||||||
let mut reader = stdout;
|
let mut reader = stdout;
|
||||||
let mut out = io::stdout();
|
let mut out = io::stdout();
|
||||||
@@ -1491,7 +1569,7 @@ pub fn run_llm_function(
|
|||||||
buf
|
buf
|
||||||
});
|
});
|
||||||
|
|
||||||
let stderr_thread = std::thread::spawn(move || {
|
let stderr_thread = thread::spawn(move || {
|
||||||
let mut buffer = [0; 1024];
|
let mut buffer = [0; 1024];
|
||||||
let mut reader = stderr;
|
let mut reader = stderr;
|
||||||
let mut err = io::stderr();
|
let mut err = io::stderr();
|
||||||
@@ -1518,9 +1596,39 @@ pub fn run_llm_function(
|
|||||||
buf
|
buf
|
||||||
});
|
});
|
||||||
|
|
||||||
let status = child
|
let timeout_secs = env::var("COYOTE_TOOL_TIMEOUT")
|
||||||
.wait()
|
.ok()
|
||||||
.map_err(|err| anyhow!("Unable to run {command_name}, {err}"))?;
|
.and_then(|v| v.parse::<u64>().ok())
|
||||||
|
.unwrap_or(1800);
|
||||||
|
let deadline = (timeout_secs > 0).then(|| Instant::now() + Duration::from_secs(timeout_secs));
|
||||||
|
let status = loop {
|
||||||
|
match child.try_wait() {
|
||||||
|
Ok(Some(status)) => break status,
|
||||||
|
Ok(None) => {}
|
||||||
|
Err(err) => bail!("Unable to run {command_name}, {err}"),
|
||||||
|
}
|
||||||
|
if let Some(deadline) = deadline
|
||||||
|
&& Instant::now() >= deadline
|
||||||
|
{
|
||||||
|
let _ = child.kill();
|
||||||
|
let _ = child.wait();
|
||||||
|
drop(stdout_thread);
|
||||||
|
drop(stderr_thread);
|
||||||
|
let tool_error_message = format!(
|
||||||
|
"Tool call '{command_name}' timed out after {timeout_secs}s and was killed (set COYOTE_TOOL_TIMEOUT to adjust; 0 = unlimited)"
|
||||||
|
);
|
||||||
|
eprintln!(
|
||||||
|
"{}",
|
||||||
|
muted_warning_text(&format!("⚠️ {tool_error_message} ⚠️"))
|
||||||
|
);
|
||||||
|
let error_json = json!({"tool_call_error": tool_error_message});
|
||||||
|
|
||||||
|
debug!("Tool call error: {error_json:?}");
|
||||||
|
|
||||||
|
return Ok(Some(error_json.to_string()));
|
||||||
|
}
|
||||||
|
thread::sleep(Duration::from_millis(100));
|
||||||
|
};
|
||||||
let stdout_bytes = stdout_thread.join().unwrap_or_default();
|
let stdout_bytes = stdout_thread.join().unwrap_or_default();
|
||||||
let stderr_bytes = stderr_thread.join().unwrap_or_default();
|
let stderr_bytes = stderr_thread.join().unwrap_or_default();
|
||||||
|
|
||||||
@@ -1540,6 +1648,11 @@ pub fn run_llm_function(
|
|||||||
if !stdout.is_empty() {
|
if !stdout.is_empty() {
|
||||||
error_json["stdout"] = json!(stdout);
|
error_json["stdout"] = json!(stdout);
|
||||||
}
|
}
|
||||||
|
if let Ok(contents) = fs::read_to_string(&tmp_file)
|
||||||
|
&& !contents.trim().is_empty()
|
||||||
|
{
|
||||||
|
error_json["output"] = json!(contents);
|
||||||
|
}
|
||||||
debug!("Tool call error: {error_json:?}");
|
debug!("Tool call error: {error_json:?}");
|
||||||
return Ok(Some(error_json.to_string()));
|
return Ok(Some(error_json.to_string()));
|
||||||
}
|
}
|
||||||
@@ -1709,7 +1822,10 @@ fn format_json_colored_keys(value: &serde_json::Value) -> String {
|
|||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
|
use crate::config::{AppState, WorkingMode};
|
||||||
|
use crate::supervisor::escalation::{EscalationQueue, EscalationRequest};
|
||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
fn call(name: &str, id: Option<&str>) -> ToolCall {
|
fn call(name: &str, id: Option<&str>) -> ToolCall {
|
||||||
ToolCall::new(name.to_string(), json!({}), id.map(|s| s.to_string()))
|
ToolCall::new(name.to_string(), json!({}), id.map(|s| s.to_string()))
|
||||||
@@ -1719,11 +1835,98 @@ mod tests {
|
|||||||
ToolCall::new(name.to_string(), args, Some("id1".to_string()))
|
ToolCall::new(name.to_string(), args, Some("id1".to_string()))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn run_async<F: Future>(f: F) -> F::Output {
|
||||||
|
tokio::runtime::Builder::new_current_thread()
|
||||||
|
.enable_all()
|
||||||
|
.build()
|
||||||
|
.unwrap()
|
||||||
|
.block_on(f)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn submit_escalation(queue: &EscalationQueue, id: &str) {
|
||||||
|
let (tx, _rx) = tokio::sync::oneshot::channel();
|
||||||
|
queue.submit(EscalationRequest {
|
||||||
|
id: id.to_string(),
|
||||||
|
from_agent_id: "a1".into(),
|
||||||
|
from_agent_name: "explore".into(),
|
||||||
|
question: "What do?".into(),
|
||||||
|
options: None,
|
||||||
|
reply_tx: tx,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn normalize_tool_result_substitutes_done_for_null() {
|
fn normalize_tool_result_substitutes_done_for_null() {
|
||||||
assert_eq!(normalize_tool_result(Value::Null), json!("DONE"));
|
assert_eq!(normalize_tool_result(Value::Null), json!("DONE"));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn inject_escalation_notification_extends_object_output() {
|
||||||
|
let mut result = ToolResult::new(call("t", Some("id-1")), json!({"status": "ok"}));
|
||||||
|
inject_escalation_notification(&mut result, vec![json!({"escalation_id": "esc_1"})]);
|
||||||
|
assert_eq!(result.output["status"], "ok");
|
||||||
|
assert_eq!(
|
||||||
|
result.output["pending_escalations"],
|
||||||
|
json!([{"escalation_id": "esc_1"}])
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
result.output["escalation_instruction"]
|
||||||
|
.as_str()
|
||||||
|
.unwrap()
|
||||||
|
.contains("agent__reply_escalation")
|
||||||
|
);
|
||||||
|
assert!(result.text.is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn inject_escalation_notification_wraps_non_object_output() {
|
||||||
|
let mut result = ToolResult::new(call("t", Some("id-1")), json!("DONE"));
|
||||||
|
inject_escalation_notification(&mut result, vec![json!({"escalation_id": "esc_2"})]);
|
||||||
|
assert_eq!(result.output["output"], json!("DONE"));
|
||||||
|
assert_eq!(
|
||||||
|
result.output["pending_escalations"],
|
||||||
|
json!([{"escalation_id": "esc_2"}])
|
||||||
|
);
|
||||||
|
assert!(result.output["escalation_instruction"].is_string());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn eval_tool_calls_soft_fails_unknown_tool() {
|
||||||
|
let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd);
|
||||||
|
let calls = vec![call("__escalation_notification", Some("id-1"))];
|
||||||
|
let results = run_async(eval_tool_calls(&mut ctx, calls)).unwrap();
|
||||||
|
assert_eq!(results.len(), 1);
|
||||||
|
let err = results[0].output["tool_call_error"].as_str().unwrap();
|
||||||
|
assert!(err.contains("Unexpected call"));
|
||||||
|
assert!(err.contains("use only tools listed in your catalog"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn eval_tool_calls_injects_escalations_into_last_result() {
|
||||||
|
let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd);
|
||||||
|
let queue = ctx.ensure_root_escalation_queue();
|
||||||
|
submit_escalation(&queue, "esc_1");
|
||||||
|
|
||||||
|
let calls = vec![call("unknown_tool", Some("id-1"))];
|
||||||
|
let results = run_async(eval_tool_calls(&mut ctx, calls)).unwrap();
|
||||||
|
|
||||||
|
assert_eq!(results.len(), 1);
|
||||||
|
assert!(
|
||||||
|
results
|
||||||
|
.iter()
|
||||||
|
.all(|r| r.call.name != "__escalation_notification")
|
||||||
|
);
|
||||||
|
let out = &results[0].output;
|
||||||
|
assert!(out["tool_call_error"].is_string());
|
||||||
|
assert_eq!(out["pending_escalations"][0]["escalation_id"], "esc_1");
|
||||||
|
assert!(
|
||||||
|
out["escalation_instruction"]
|
||||||
|
.as_str()
|
||||||
|
.unwrap()
|
||||||
|
.contains("agent__reply_escalation")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn normalize_tool_result_preserves_non_null_values() {
|
fn normalize_tool_result_preserves_non_null_values() {
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
@@ -2151,4 +2354,149 @@ mod tests {
|
|||||||
let tc = call_with_args("t", json!(42));
|
let tc = call_with_args("t", json!(42));
|
||||||
assert!(tc.parse_arguments().is_err());
|
assert!(tc.parse_arguments().is_err());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn write_file_atomic_writes_and_skips_unchanged() {
|
||||||
|
let dir = temp_file("-atomic-", "");
|
||||||
|
fs::create_dir_all(&dir).unwrap();
|
||||||
|
let path = dir.join("shim");
|
||||||
|
|
||||||
|
write_file_atomic(&path, "one", Some(0o755)).unwrap();
|
||||||
|
assert_eq!(fs::read_to_string(&path).unwrap(), "one");
|
||||||
|
#[cfg(unix)]
|
||||||
|
{
|
||||||
|
use std::os::unix::fs::PermissionsExt;
|
||||||
|
assert_eq!(
|
||||||
|
fs::metadata(&path).unwrap().permissions().mode() & 0o777,
|
||||||
|
0o755
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(unix)]
|
||||||
|
{
|
||||||
|
use std::os::unix::fs::MetadataExt;
|
||||||
|
let ino = fs::metadata(&path).unwrap().ino();
|
||||||
|
write_file_atomic(&path, "one", Some(0o755)).unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
fs::metadata(&path).unwrap().ino(),
|
||||||
|
ino,
|
||||||
|
"unchanged content must not be rewritten"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
write_file_atomic(&path, "two", None).unwrap();
|
||||||
|
assert_eq!(fs::read_to_string(&path).unwrap(), "two");
|
||||||
|
assert_eq!(
|
||||||
|
fs::read_dir(&dir).unwrap().count(),
|
||||||
|
1,
|
||||||
|
"no tmp files left behind"
|
||||||
|
);
|
||||||
|
|
||||||
|
fs::remove_dir_all(&dir).unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn write_file_atomic_concurrent_writers_to_same_target() {
|
||||||
|
let dir = temp_file("-atomic-concurrent-", "");
|
||||||
|
fs::create_dir_all(&dir).unwrap();
|
||||||
|
let path = dir.join("shim");
|
||||||
|
|
||||||
|
let contents: Vec<String> = (0..8)
|
||||||
|
.map(|i| format!("#!/bin/sh\necho writer-{i}\n"))
|
||||||
|
.collect();
|
||||||
|
thread::scope(|scope| {
|
||||||
|
for content in &contents {
|
||||||
|
scope.spawn(|| {
|
||||||
|
for _ in 0..50 {
|
||||||
|
write_file_atomic(&path, content, Some(0o755)).unwrap();
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
let final_content = fs::read_to_string(&path).unwrap();
|
||||||
|
assert!(
|
||||||
|
contents.contains(&final_content),
|
||||||
|
"final content must be one writer's complete content, got: {final_content:?}"
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
fs::read_dir(&dir).unwrap().count(),
|
||||||
|
1,
|
||||||
|
"no tmp files left behind"
|
||||||
|
);
|
||||||
|
|
||||||
|
fs::remove_dir_all(&dir).unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(unix)]
|
||||||
|
#[test]
|
||||||
|
fn run_llm_function_includes_llm_output_on_nonzero_exit() {
|
||||||
|
let result = run_llm_function(
|
||||||
|
"bash".into(),
|
||||||
|
vec![
|
||||||
|
"-c".into(),
|
||||||
|
"echo partial-output >> \"$LLM_OUTPUT\"; echo err-text >&2; exit 3".into(),
|
||||||
|
],
|
||||||
|
HashMap::new(),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.unwrap()
|
||||||
|
.expect("nonzero exit must return an error payload");
|
||||||
|
|
||||||
|
let json: serde_json::Value = serde_json::from_str(&result).unwrap();
|
||||||
|
assert!(
|
||||||
|
json["tool_call_error"]
|
||||||
|
.as_str()
|
||||||
|
.unwrap()
|
||||||
|
.contains("exited with code 3")
|
||||||
|
);
|
||||||
|
assert_eq!(json["stderr"], "err-text");
|
||||||
|
assert_eq!(json["output"], "partial-output\n");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn bin_entry_stem_strips_run_prefix_and_extension() {
|
||||||
|
assert_eq!(bin_entry_stem("fs_grep"), "fs_grep");
|
||||||
|
assert_eq!(bin_entry_stem("fs_grep.cmd"), "fs_grep");
|
||||||
|
assert_eq!(bin_entry_stem("run-web_search.ts"), "web_search");
|
||||||
|
assert_eq!(bin_entry_stem("run-fs_grep.sh"), "fs_grep");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn prune_stale_bin_entries_removes_only_stale_files() {
|
||||||
|
let dir = temp_file("-prune-", "");
|
||||||
|
fs::create_dir_all(dir.join("nested")).unwrap();
|
||||||
|
for name in [
|
||||||
|
"fs_grep",
|
||||||
|
"run-web_search.ts",
|
||||||
|
"old_tool",
|
||||||
|
"run-old_tool.ts",
|
||||||
|
"myagent",
|
||||||
|
] {
|
||||||
|
fs::write(dir.join(name), "x").unwrap();
|
||||||
|
}
|
||||||
|
let valid_stems: HashSet<String> = ["fs_grep", "web_search"]
|
||||||
|
.iter()
|
||||||
|
.map(|s| s.to_string())
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
prune_stale_bin_entries(&dir, &valid_stems, Some("myagent")).unwrap();
|
||||||
|
|
||||||
|
assert!(dir.join("fs_grep").exists());
|
||||||
|
assert!(dir.join("run-web_search.ts").exists());
|
||||||
|
assert!(dir.join("myagent").exists(), "agent binary must survive");
|
||||||
|
assert!(!dir.join("old_tool").exists());
|
||||||
|
assert!(!dir.join("run-old_tool.ts").exists());
|
||||||
|
assert!(!dir.join("nested").exists(), "directories must be removed");
|
||||||
|
|
||||||
|
fs::remove_dir_all(&dir).unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn prune_stale_bin_entries_creates_missing_dir() {
|
||||||
|
let dir = temp_file("-prune-missing-", "");
|
||||||
|
prune_stale_bin_entries(&dir, &HashSet::new(), None).unwrap();
|
||||||
|
assert!(dir.is_dir());
|
||||||
|
fs::remove_dir_all(&dir).unwrap();
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+109
-1
@@ -85,7 +85,17 @@ pub fn check_pending_agents_guardrail(ctx: &mut RequestContext) -> GuardrailActi
|
|||||||
}
|
}
|
||||||
|
|
||||||
ctx.pending_agents_guardrail_count += 1;
|
ctx.pending_agents_guardrail_count += 1;
|
||||||
GuardrailAction::Inject(build_pending_agents_guardrail_prompt(&pending))
|
let mut prompt = build_pending_agents_guardrail_prompt(&pending);
|
||||||
|
if let Some(queue) = ctx.root_escalation_queue()
|
||||||
|
&& queue.has_pending()
|
||||||
|
{
|
||||||
|
let summary = serde_json::to_string(&queue.pending_summary()).unwrap_or_default();
|
||||||
|
prompt.push_str(&format!(
|
||||||
|
"\n\nAdditionally, child agents have pending escalations blocking them. Reply to each \
|
||||||
|
via `agent__reply_escalation` first:\n{summary}"
|
||||||
|
));
|
||||||
|
}
|
||||||
|
GuardrailAction::Inject(prompt)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn escalation_function_declarations() -> Vec<FunctionDeclaration> {
|
pub fn escalation_function_declarations() -> Vec<FunctionDeclaration> {
|
||||||
@@ -1996,4 +2006,102 @@ mod tests {
|
|||||||
let q2 = ctx.ensure_root_escalation_queue();
|
let q2 = ctx.ensure_root_escalation_queue();
|
||||||
assert!(Arc::ptr_eq(&q1, &q2));
|
assert!(Arc::ptr_eq(&q1, &q2));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn guardrail_prompt_mentions_pending_escalations() {
|
||||||
|
let rt = tokio::runtime::Builder::new_current_thread()
|
||||||
|
.enable_all()
|
||||||
|
.build()
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
rt.block_on(async {
|
||||||
|
let mut ctx = ctx_with_supervisor(4, 3);
|
||||||
|
let join_handle = tokio::spawn(async {
|
||||||
|
time::sleep(Duration::from_secs(60)).await;
|
||||||
|
Ok(AgentResult {
|
||||||
|
id: "slow".into(),
|
||||||
|
agent_name: "test".into(),
|
||||||
|
output: String::new(),
|
||||||
|
exit_status: AgentExitStatus::Completed,
|
||||||
|
})
|
||||||
|
});
|
||||||
|
let handle = AgentHandle {
|
||||||
|
id: "slow".into(),
|
||||||
|
agent_name: "test".into(),
|
||||||
|
depth: 1,
|
||||||
|
inbox: Arc::new(Inbox::new()),
|
||||||
|
abort_signal: create_abort_signal(),
|
||||||
|
join_handle,
|
||||||
|
child_supervisor: None,
|
||||||
|
};
|
||||||
|
ctx.supervisor
|
||||||
|
.as_ref()
|
||||||
|
.unwrap()
|
||||||
|
.write()
|
||||||
|
.register(handle)
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let queue = ctx.ensure_root_escalation_queue();
|
||||||
|
let (tx, _rx) = tokio::sync::oneshot::channel();
|
||||||
|
queue.submit(EscalationRequest {
|
||||||
|
id: "esc_9".into(),
|
||||||
|
from_agent_id: "a1".into(),
|
||||||
|
from_agent_name: "explore".into(),
|
||||||
|
question: "Which option?".into(),
|
||||||
|
options: None,
|
||||||
|
reply_tx: tx,
|
||||||
|
});
|
||||||
|
|
||||||
|
match check_pending_agents_guardrail(&mut ctx) {
|
||||||
|
GuardrailAction::Inject(prompt) => {
|
||||||
|
assert!(prompt.contains("agent__reply_escalation"));
|
||||||
|
assert!(prompt.contains("esc_9"));
|
||||||
|
}
|
||||||
|
_ => panic!("expected Inject action"),
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn guardrail_prompt_omits_escalations_when_none_pending() {
|
||||||
|
let rt = tokio::runtime::Builder::new_current_thread()
|
||||||
|
.enable_all()
|
||||||
|
.build()
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
rt.block_on(async {
|
||||||
|
let mut ctx = ctx_with_supervisor(4, 3);
|
||||||
|
let join_handle = tokio::spawn(async {
|
||||||
|
time::sleep(Duration::from_secs(60)).await;
|
||||||
|
Ok(AgentResult {
|
||||||
|
id: "slow".into(),
|
||||||
|
agent_name: "test".into(),
|
||||||
|
output: String::new(),
|
||||||
|
exit_status: AgentExitStatus::Completed,
|
||||||
|
})
|
||||||
|
});
|
||||||
|
let handle = AgentHandle {
|
||||||
|
id: "slow".into(),
|
||||||
|
agent_name: "test".into(),
|
||||||
|
depth: 1,
|
||||||
|
inbox: Arc::new(Inbox::new()),
|
||||||
|
abort_signal: create_abort_signal(),
|
||||||
|
join_handle,
|
||||||
|
child_supervisor: None,
|
||||||
|
};
|
||||||
|
ctx.supervisor
|
||||||
|
.as_ref()
|
||||||
|
.unwrap()
|
||||||
|
.write()
|
||||||
|
.register(handle)
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
match check_pending_agents_guardrail(&mut ctx) {
|
||||||
|
GuardrailAction::Inject(prompt) => {
|
||||||
|
assert!(!prompt.contains("agent__reply_escalation"));
|
||||||
|
}
|
||||||
|
_ => panic!("expected Inject action"),
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+64
-1
@@ -250,7 +250,30 @@ impl GraphExecutor {
|
|||||||
branch_tasks.push(task);
|
branch_tasks.push(task);
|
||||||
}
|
}
|
||||||
|
|
||||||
let joined = join_all(branch_tasks).await;
|
let joined = match graph_timeout {
|
||||||
|
Some(t) => {
|
||||||
|
let remaining = t.saturating_sub(start.elapsed());
|
||||||
|
let abort_handles: Vec<_> = branch_tasks
|
||||||
|
.iter()
|
||||||
|
.map(|task| task.abort_handle())
|
||||||
|
.collect();
|
||||||
|
match tokio::time::timeout(remaining, join_all(branch_tasks)).await {
|
||||||
|
Ok(joined) => joined,
|
||||||
|
Err(_) => {
|
||||||
|
for handle in abort_handles {
|
||||||
|
handle.abort();
|
||||||
|
}
|
||||||
|
bail!(
|
||||||
|
"Graph '{}' timed out after {}s during super-step with frontier {:?}",
|
||||||
|
graph.name,
|
||||||
|
t.as_secs(),
|
||||||
|
sorted_frontier(&frontier)
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
None => join_all(branch_tasks).await,
|
||||||
|
};
|
||||||
|
|
||||||
let mut branch_writes: Vec<BranchWrites> = Vec::new();
|
let mut branch_writes: Vec<BranchWrites> = Vec::new();
|
||||||
let mut next_frontier: HashSet<String> = HashSet::new();
|
let mut next_frontier: HashSet<String> = HashSet::new();
|
||||||
@@ -793,4 +816,44 @@ nodes:
|
|||||||
"error should list both End nodes: {err}"
|
"error should list both End nodes: {err}"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn graph_timeout_interrupts_in_flight_super_step() {
|
||||||
|
if !cmd_available("bash") {
|
||||||
|
eprintln!("skipping: bash not available");
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let ws = TestWorkspace::new();
|
||||||
|
ws.write_script("sleeper.sh", "#!/bin/bash\nsleep 10\necho '{}'\n");
|
||||||
|
|
||||||
|
let yaml = r#"
|
||||||
|
name: inflight_timeout_test
|
||||||
|
start: sleeper
|
||||||
|
settings:
|
||||||
|
timeout: 1
|
||||||
|
nodes:
|
||||||
|
sleeper:
|
||||||
|
type: script
|
||||||
|
script: sleeper.sh
|
||||||
|
state_updates: {}
|
||||||
|
next: done
|
||||||
|
done:
|
||||||
|
type: end
|
||||||
|
output: "done"
|
||||||
|
"#;
|
||||||
|
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
||||||
|
let mut ctx = make_ctx();
|
||||||
|
let abort = create_abort_signal();
|
||||||
|
let result = GraphExecutor::new(graph, &ws.dir)
|
||||||
|
.execute(&mut ctx, abort)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
assert!(result.is_err(), "expected in-flight timeout to error");
|
||||||
|
let err = format!("{:#}", result.unwrap_err());
|
||||||
|
assert!(
|
||||||
|
err.contains("timed out after 1s during super-step"),
|
||||||
|
"error should report during-super-step timeout: {err}"
|
||||||
|
);
|
||||||
|
assert!(err.contains("sleeper"), "error should name frontier: {err}");
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -54,6 +54,7 @@ impl ScriptExecutor {
|
|||||||
let mut cmd = build_command(language, &script_path)?;
|
let mut cmd = build_command(language, &script_path)?;
|
||||||
cmd.stdout(Stdio::piped());
|
cmd.stdout(Stdio::piped());
|
||||||
cmd.stderr(Stdio::piped());
|
cmd.stderr(Stdio::piped());
|
||||||
|
cmd.kill_on_drop(true);
|
||||||
cmd.envs(&self.extra_envs);
|
cmd.envs(&self.extra_envs);
|
||||||
cmd.env("AUTO_CONFIRM", "true");
|
cmd.env("AUTO_CONFIRM", "true");
|
||||||
match &state_repr {
|
match &state_repr {
|
||||||
|
|||||||
@@ -0,0 +1,651 @@
|
|||||||
|
use crate::mcp::oauth::{force_refresh_mcp_token, load_or_refresh_mcp_token};
|
||||||
|
use http::{HeaderName, HeaderValue};
|
||||||
|
use log::debug;
|
||||||
|
use rmcp::model::ClientJsonRpcMessage;
|
||||||
|
use rmcp::transport::common::client_side_sse::BoxedSseResponse;
|
||||||
|
use rmcp::transport::streamable_http_client::{
|
||||||
|
AuthRequiredError, StreamableHttpClient, StreamableHttpError, StreamableHttpPostResponse,
|
||||||
|
};
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::future::Future;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
/// [`StreamableHttpClient`] wrapper that injects the OAuth bearer token for an
|
||||||
|
/// MCP server on every request instead of pinning it at spawn time, so tokens
|
||||||
|
/// refreshed mid-session take effect without reconnecting.
|
||||||
|
///
|
||||||
|
/// A caller-supplied `auth_header` always passes through untouched; only a
|
||||||
|
/// `None` header is filled from the stored token. When the wrapper injected
|
||||||
|
/// the token and a POST comes back 401, it forces a token refresh and retries
|
||||||
|
/// exactly once (see [`Self::post_with_retry`]).
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct McpOAuthClient<C = reqwest::Client> {
|
||||||
|
inner: C,
|
||||||
|
server: Arc<str>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<C> McpOAuthClient<C> {
|
||||||
|
pub fn new(inner: C, server: &str) -> Self {
|
||||||
|
Self {
|
||||||
|
inner,
|
||||||
|
server: Arc::from(server),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<C: StreamableHttpClient + Sync> McpOAuthClient<C> {
|
||||||
|
/// Resolves the effective auth header. Caller-supplied values pass through
|
||||||
|
/// untouched; `None` is filled from the stored token for this server.
|
||||||
|
/// Returns the header plus whether the wrapper injected it. Errors with
|
||||||
|
/// [`StreamableHttpError::AuthRequired`] (without contacting the server)
|
||||||
|
/// when no usable token exists.
|
||||||
|
async fn resolve_auth(
|
||||||
|
&self,
|
||||||
|
auth_header: Option<String>,
|
||||||
|
) -> Result<(Option<String>, bool), StreamableHttpError<C::Error>> {
|
||||||
|
if auth_header.is_some() {
|
||||||
|
return Ok((auth_header, false));
|
||||||
|
}
|
||||||
|
match load_or_refresh_mcp_token(&self.server).await.into_token() {
|
||||||
|
Some(token) => Ok((Some(token), true)),
|
||||||
|
None => Err(self.auth_required()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn auth_required(&self) -> StreamableHttpError<C::Error> {
|
||||||
|
StreamableHttpError::AuthRequired(AuthRequiredError::new(format!(
|
||||||
|
"no valid OAuth token for MCP server '{server}'; \
|
||||||
|
run `.mcp auth {server}` to re-authenticate",
|
||||||
|
server = self.server
|
||||||
|
)))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn post_with_retry<F, Fut>(
|
||||||
|
&self,
|
||||||
|
auth_header: Option<String>,
|
||||||
|
mut post: F,
|
||||||
|
) -> Result<StreamableHttpPostResponse, StreamableHttpError<C::Error>>
|
||||||
|
where
|
||||||
|
F: FnMut(Option<String>) -> Fut,
|
||||||
|
Fut: Future<Output = Result<StreamableHttpPostResponse, StreamableHttpError<C::Error>>>,
|
||||||
|
{
|
||||||
|
let (auth, injected) = self.resolve_auth(auth_header).await?;
|
||||||
|
let (original, rejected) = match (post(auth.clone()).await, auth) {
|
||||||
|
(Err(err @ StreamableHttpError::AuthRequired(_)), Some(rejected)) if injected => {
|
||||||
|
(err, rejected)
|
||||||
|
}
|
||||||
|
(result, _) => return result,
|
||||||
|
};
|
||||||
|
|
||||||
|
debug!(
|
||||||
|
"MCP server '{}' rejected the injected token; forcing a refresh and retrying once",
|
||||||
|
self.server
|
||||||
|
);
|
||||||
|
|
||||||
|
let Some(token) = force_refresh_mcp_token(&self.server, &rejected).await else {
|
||||||
|
return Err(original);
|
||||||
|
};
|
||||||
|
|
||||||
|
match post(Some(token)).await {
|
||||||
|
Err(StreamableHttpError::AuthRequired(_)) => {
|
||||||
|
debug!(
|
||||||
|
"Retry after forced token refresh was rejected again by MCP server '{}'",
|
||||||
|
self.server
|
||||||
|
);
|
||||||
|
Err(original)
|
||||||
|
}
|
||||||
|
result => result,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<C: StreamableHttpClient + Sync> StreamableHttpClient for McpOAuthClient<C> {
|
||||||
|
type Error = C::Error;
|
||||||
|
|
||||||
|
async fn post_message(
|
||||||
|
&self,
|
||||||
|
uri: Arc<str>,
|
||||||
|
message: ClientJsonRpcMessage,
|
||||||
|
session_id: Option<Arc<str>>,
|
||||||
|
auth_header: Option<String>,
|
||||||
|
custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||||
|
) -> Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>> {
|
||||||
|
self.post_with_retry(auth_header, |auth| {
|
||||||
|
self.inner.post_message(
|
||||||
|
uri.clone(),
|
||||||
|
message.clone(),
|
||||||
|
session_id.clone(),
|
||||||
|
auth,
|
||||||
|
custom_headers.clone(),
|
||||||
|
)
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Overridden rather than left to the trait default: the default impl
|
||||||
|
/// delegates to [`Self::post_message`], silently dropping the
|
||||||
|
/// transport-wide SSE event size limit. Delegating to the inner client's
|
||||||
|
/// size-enforcing variant keeps the limit applied at the raw byte layer.
|
||||||
|
async fn post_message_with_max_sse_event_size(
|
||||||
|
&self,
|
||||||
|
uri: Arc<str>,
|
||||||
|
message: ClientJsonRpcMessage,
|
||||||
|
session_id: Option<Arc<str>>,
|
||||||
|
auth_header: Option<String>,
|
||||||
|
custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||||
|
max_sse_event_size: usize,
|
||||||
|
) -> Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>> {
|
||||||
|
self.post_with_retry(auth_header, |auth| {
|
||||||
|
self.inner.post_message_with_max_sse_event_size(
|
||||||
|
uri.clone(),
|
||||||
|
message.clone(),
|
||||||
|
session_id.clone(),
|
||||||
|
auth,
|
||||||
|
custom_headers.clone(),
|
||||||
|
max_sse_event_size,
|
||||||
|
)
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn delete_session(
|
||||||
|
&self,
|
||||||
|
uri: Arc<str>,
|
||||||
|
session_id: Arc<str>,
|
||||||
|
auth_header: Option<String>,
|
||||||
|
custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||||
|
) -> Result<(), StreamableHttpError<Self::Error>> {
|
||||||
|
let (auth, _) = self.resolve_auth(auth_header).await?;
|
||||||
|
self.inner
|
||||||
|
.delete_session(uri, session_id, auth, custom_headers)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn get_stream(
|
||||||
|
&self,
|
||||||
|
uri: Arc<str>,
|
||||||
|
session_id: Option<Arc<str>>,
|
||||||
|
last_event_id: Option<String>,
|
||||||
|
auth_header: Option<String>,
|
||||||
|
custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||||
|
) -> Result<BoxedSseResponse, StreamableHttpError<Self::Error>> {
|
||||||
|
let (auth, _) = self.resolve_auth(auth_header).await?;
|
||||||
|
self.inner
|
||||||
|
.get_stream(uri, session_id, last_event_id, auth, custom_headers)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Overridden for the same reason as
|
||||||
|
/// [`Self::post_message_with_max_sse_event_size`]: the trait default
|
||||||
|
/// bypasses SSE event size enforcement.
|
||||||
|
async fn get_stream_with_max_sse_event_size(
|
||||||
|
&self,
|
||||||
|
uri: Arc<str>,
|
||||||
|
session_id: Option<Arc<str>>,
|
||||||
|
last_event_id: Option<String>,
|
||||||
|
auth_header: Option<String>,
|
||||||
|
custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||||
|
max_sse_event_size: usize,
|
||||||
|
) -> Result<BoxedSseResponse, StreamableHttpError<Self::Error>> {
|
||||||
|
let (auth, _) = self.resolve_auth(auth_header).await?;
|
||||||
|
self.inner
|
||||||
|
.get_stream_with_max_sse_event_size(
|
||||||
|
uri,
|
||||||
|
session_id,
|
||||||
|
last_event_id,
|
||||||
|
auth,
|
||||||
|
custom_headers,
|
||||||
|
max_sse_event_size,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use crate::config::paths;
|
||||||
|
use crate::mcp::oauth::test_support::with_temp_cache;
|
||||||
|
use futures_util::StreamExt;
|
||||||
|
use parking_lot::Mutex;
|
||||||
|
use serial_test::serial;
|
||||||
|
use std::convert::Infallible;
|
||||||
|
use std::fs;
|
||||||
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||||
|
|
||||||
|
const FRESH: i64 = 9999999999;
|
||||||
|
|
||||||
|
type Calls = Arc<Mutex<Vec<(&'static str, Option<String>)>>>;
|
||||||
|
|
||||||
|
type AfterFirstCallAction = Option<Box<dyn FnOnce() + Send + 'static>>;
|
||||||
|
|
||||||
|
/// Inner client that records `(method, auth_header)` per call, rejects the
|
||||||
|
/// first `reject_times` POSTs/streams with `AuthRequired` (then the next
|
||||||
|
/// `transport_error_times` with a non-auth error), and runs an optional
|
||||||
|
/// side effect after the first call (to mutate token files between the
|
||||||
|
/// initial attempt and the retry).
|
||||||
|
#[derive(Clone, Default)]
|
||||||
|
struct FakeInner {
|
||||||
|
calls: Calls,
|
||||||
|
reject_times: Arc<AtomicUsize>,
|
||||||
|
transport_error_times: Arc<AtomicUsize>,
|
||||||
|
after_first_call: Arc<Mutex<AfterFirstCallAction>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl FakeInner {
|
||||||
|
fn record(
|
||||||
|
&self,
|
||||||
|
method: &'static str,
|
||||||
|
auth: Option<String>,
|
||||||
|
) -> Result<(), StreamableHttpError<Infallible>> {
|
||||||
|
self.calls.lock().push((method, auth));
|
||||||
|
if let Some(f) = self.after_first_call.lock().take() {
|
||||||
|
f();
|
||||||
|
}
|
||||||
|
if self.reject_times.load(Ordering::SeqCst) > 0 {
|
||||||
|
self.reject_times.fetch_sub(1, Ordering::SeqCst);
|
||||||
|
return Err(Self::rejection());
|
||||||
|
}
|
||||||
|
if self.transport_error_times.load(Ordering::SeqCst) > 0 {
|
||||||
|
self.transport_error_times.fetch_sub(1, Ordering::SeqCst);
|
||||||
|
return Err(StreamableHttpError::UnexpectedServerResponse(
|
||||||
|
"connection reset".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn rejection() -> StreamableHttpError<Infallible> {
|
||||||
|
StreamableHttpError::AuthRequired(AuthRequiredError::new(
|
||||||
|
"Bearer error=\"invalid_token\"".to_string(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl StreamableHttpClient for FakeInner {
|
||||||
|
type Error = Infallible;
|
||||||
|
|
||||||
|
async fn post_message(
|
||||||
|
&self,
|
||||||
|
_uri: Arc<str>,
|
||||||
|
_message: ClientJsonRpcMessage,
|
||||||
|
_session_id: Option<Arc<str>>,
|
||||||
|
auth_header: Option<String>,
|
||||||
|
_custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||||
|
) -> Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>> {
|
||||||
|
self.record("post_message", auth_header)?;
|
||||||
|
Ok(StreamableHttpPostResponse::Accepted)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn post_message_with_max_sse_event_size(
|
||||||
|
&self,
|
||||||
|
_uri: Arc<str>,
|
||||||
|
_message: ClientJsonRpcMessage,
|
||||||
|
_session_id: Option<Arc<str>>,
|
||||||
|
auth_header: Option<String>,
|
||||||
|
_custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||||
|
_max_sse_event_size: usize,
|
||||||
|
) -> Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>> {
|
||||||
|
self.record("post_message_with_max_sse_event_size", auth_header)?;
|
||||||
|
Ok(StreamableHttpPostResponse::Accepted)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn delete_session(
|
||||||
|
&self,
|
||||||
|
_uri: Arc<str>,
|
||||||
|
_session_id: Arc<str>,
|
||||||
|
auth_header: Option<String>,
|
||||||
|
_custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||||
|
) -> Result<(), StreamableHttpError<Self::Error>> {
|
||||||
|
self.record("delete_session", auth_header)?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn get_stream(
|
||||||
|
&self,
|
||||||
|
_uri: Arc<str>,
|
||||||
|
_session_id: Option<Arc<str>>,
|
||||||
|
_last_event_id: Option<String>,
|
||||||
|
auth_header: Option<String>,
|
||||||
|
_custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||||
|
) -> Result<BoxedSseResponse, StreamableHttpError<Self::Error>> {
|
||||||
|
self.record("get_stream", auth_header)?;
|
||||||
|
Ok(futures_util::stream::empty().boxed())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn get_stream_with_max_sse_event_size(
|
||||||
|
&self,
|
||||||
|
_uri: Arc<str>,
|
||||||
|
_session_id: Option<Arc<str>>,
|
||||||
|
_last_event_id: Option<String>,
|
||||||
|
auth_header: Option<String>,
|
||||||
|
_custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||||
|
_max_sse_event_size: usize,
|
||||||
|
) -> Result<BoxedSseResponse, StreamableHttpError<Self::Error>> {
|
||||||
|
self.record("get_stream_with_max_sse_event_size", auth_header)?;
|
||||||
|
Ok(futures_util::stream::empty().boxed())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn write_token_file(server: &str, access_token: &str, expires_at: i64) {
|
||||||
|
fs::create_dir_all(paths::oauth_tokens_dir()).unwrap();
|
||||||
|
fs::write(
|
||||||
|
paths::token_file(&format!("mcp_{server}")),
|
||||||
|
format!(
|
||||||
|
r#"{{"access_token":"{access_token}","refresh_token":"r","expires_at":{expires_at}}}"#
|
||||||
|
),
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
|
fn ping() -> ClientJsonRpcMessage {
|
||||||
|
serde_json::from_value(serde_json::json!({
|
||||||
|
"jsonrpc": "2.0",
|
||||||
|
"id": 1,
|
||||||
|
"method": "ping"
|
||||||
|
}))
|
||||||
|
.unwrap()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn rt() -> tokio::runtime::Runtime {
|
||||||
|
tokio::runtime::Builder::new_current_thread()
|
||||||
|
.enable_all()
|
||||||
|
.build()
|
||||||
|
.unwrap()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn post(
|
||||||
|
client: &McpOAuthClient<FakeInner>,
|
||||||
|
auth_header: Option<String>,
|
||||||
|
) -> Result<StreamableHttpPostResponse, StreamableHttpError<Infallible>> {
|
||||||
|
rt().block_on(client.post_message(
|
||||||
|
Arc::from("http://mcp.test/mcp"),
|
||||||
|
ping(),
|
||||||
|
None,
|
||||||
|
auth_header,
|
||||||
|
HashMap::new(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[serial]
|
||||||
|
fn injects_token_from_disk_when_auth_header_none() {
|
||||||
|
with_temp_cache(|| {
|
||||||
|
write_token_file("wrapper-inject", "tok-live", FRESH);
|
||||||
|
let inner = FakeInner::default();
|
||||||
|
let client = McpOAuthClient::new(inner.clone(), "wrapper-inject");
|
||||||
|
|
||||||
|
let result = post(&client, None);
|
||||||
|
|
||||||
|
assert!(result.is_ok());
|
||||||
|
assert_eq!(
|
||||||
|
*inner.calls.lock(),
|
||||||
|
vec![("post_message", Some("tok-live".to_string()))]
|
||||||
|
);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[serial]
|
||||||
|
fn caller_supplied_auth_header_passes_through() {
|
||||||
|
with_temp_cache(|| {
|
||||||
|
write_token_file("wrapper-passthrough", "tok-disk", FRESH);
|
||||||
|
let inner = FakeInner::default();
|
||||||
|
let client = McpOAuthClient::new(inner.clone(), "wrapper-passthrough");
|
||||||
|
|
||||||
|
let result = post(&client, Some("caller-tok".to_string()));
|
||||||
|
|
||||||
|
assert!(result.is_ok());
|
||||||
|
assert_eq!(
|
||||||
|
*inner.calls.lock(),
|
||||||
|
vec![("post_message", Some("caller-tok".to_string()))]
|
||||||
|
);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[serial]
|
||||||
|
fn caller_supplied_header_rejection_propagates_without_refresh() {
|
||||||
|
with_temp_cache(|| {
|
||||||
|
// A fresh, different token sits on disk: if the injected guard
|
||||||
|
// were dropped, the wrapper would refresh and retry with it.
|
||||||
|
write_token_file("wrapper-caller-401", "tok-disk", FRESH);
|
||||||
|
let inner = FakeInner::default();
|
||||||
|
inner.reject_times.store(1, Ordering::SeqCst);
|
||||||
|
let client = McpOAuthClient::new(inner.clone(), "wrapper-caller-401");
|
||||||
|
|
||||||
|
let result = post(&client, Some("caller-tok".to_string()));
|
||||||
|
|
||||||
|
assert!(matches!(result, Err(StreamableHttpError::AuthRequired(_))));
|
||||||
|
assert_eq!(
|
||||||
|
*inner.calls.lock(),
|
||||||
|
vec![("post_message", Some("caller-tok".to_string()))]
|
||||||
|
);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[serial]
|
||||||
|
fn missing_token_returns_auth_required_without_calling_inner() {
|
||||||
|
with_temp_cache(|| {
|
||||||
|
let inner = FakeInner::default();
|
||||||
|
let client = McpOAuthClient::new(inner.clone(), "wrapper-no-token");
|
||||||
|
|
||||||
|
let result = post(&client, None);
|
||||||
|
|
||||||
|
assert!(matches!(result, Err(StreamableHttpError::AuthRequired(_))));
|
||||||
|
assert!(inner.calls.lock().is_empty());
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[serial]
|
||||||
|
fn rejected_token_forces_refresh_and_retries_once() {
|
||||||
|
with_temp_cache(|| {
|
||||||
|
write_token_file("wrapper-retry", "tok-a", FRESH);
|
||||||
|
let inner = FakeInner::default();
|
||||||
|
inner.reject_times.store(1, Ordering::SeqCst);
|
||||||
|
// Simulate a concurrent refresh landing between the rejection and
|
||||||
|
// the forced refresh: the retry must carry the new token.
|
||||||
|
*inner.after_first_call.lock() = Some(Box::new(|| {
|
||||||
|
write_token_file("wrapper-retry", "tok-b", FRESH);
|
||||||
|
}));
|
||||||
|
let client = McpOAuthClient::new(inner.clone(), "wrapper-retry");
|
||||||
|
|
||||||
|
let result = post(&client, None);
|
||||||
|
|
||||||
|
assert!(result.is_ok());
|
||||||
|
assert_eq!(
|
||||||
|
*inner.calls.lock(),
|
||||||
|
vec![
|
||||||
|
("post_message", Some("tok-a".to_string())),
|
||||||
|
("post_message", Some("tok-b".to_string())),
|
||||||
|
]
|
||||||
|
);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[serial]
|
||||||
|
fn failed_force_refresh_propagates_original_error_after_one_call() {
|
||||||
|
with_temp_cache(|| {
|
||||||
|
write_token_file("wrapper-refresh-fail", "tok-a", FRESH);
|
||||||
|
let inner = FakeInner::default();
|
||||||
|
inner.reject_times.store(1, Ordering::SeqCst);
|
||||||
|
// Token file gone by refresh time: force refresh yields nothing.
|
||||||
|
*inner.after_first_call.lock() = Some(Box::new(|| {
|
||||||
|
fs::remove_file(paths::token_file("mcp_wrapper-refresh-fail")).unwrap();
|
||||||
|
}));
|
||||||
|
let client = McpOAuthClient::new(inner.clone(), "wrapper-refresh-fail");
|
||||||
|
|
||||||
|
let result = post(&client, None);
|
||||||
|
|
||||||
|
assert!(matches!(result, Err(StreamableHttpError::AuthRequired(_))));
|
||||||
|
assert_eq!(inner.calls.lock().len(), 1);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[serial]
|
||||||
|
fn second_rejection_after_retry_propagates_original_error() {
|
||||||
|
with_temp_cache(|| {
|
||||||
|
write_token_file("wrapper-double-401", "tok-a", FRESH);
|
||||||
|
let inner = FakeInner::default();
|
||||||
|
inner.reject_times.store(2, Ordering::SeqCst);
|
||||||
|
// A changed token appears before the forced refresh, so the retry
|
||||||
|
// actually runs (an unchanged token would trigger a real refresh
|
||||||
|
// attempt, which fails without a cached registration).
|
||||||
|
*inner.after_first_call.lock() = Some(Box::new(|| {
|
||||||
|
write_token_file("wrapper-double-401", "tok-b", FRESH);
|
||||||
|
}));
|
||||||
|
let client = McpOAuthClient::new(inner.clone(), "wrapper-double-401");
|
||||||
|
|
||||||
|
let result = post(&client, None);
|
||||||
|
|
||||||
|
assert!(matches!(result, Err(StreamableHttpError::AuthRequired(_))));
|
||||||
|
assert_eq!(inner.calls.lock().len(), 2);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[serial]
|
||||||
|
fn non_auth_retry_error_propagates_as_is() {
|
||||||
|
with_temp_cache(|| {
|
||||||
|
write_token_file("wrapper-retry-transport", "tok-a", FRESH);
|
||||||
|
let inner = FakeInner::default();
|
||||||
|
inner.reject_times.store(1, Ordering::SeqCst);
|
||||||
|
inner.transport_error_times.store(1, Ordering::SeqCst);
|
||||||
|
*inner.after_first_call.lock() = Some(Box::new(|| {
|
||||||
|
write_token_file("wrapper-retry-transport", "tok-b", FRESH);
|
||||||
|
}));
|
||||||
|
let client = McpOAuthClient::new(inner.clone(), "wrapper-retry-transport");
|
||||||
|
|
||||||
|
let result = post(&client, None);
|
||||||
|
|
||||||
|
assert!(matches!(
|
||||||
|
result,
|
||||||
|
Err(StreamableHttpError::UnexpectedServerResponse(_))
|
||||||
|
));
|
||||||
|
assert_eq!(inner.calls.lock().len(), 2);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[serial]
|
||||||
|
fn sized_post_delegates_to_inner_sized_variant() {
|
||||||
|
with_temp_cache(|| {
|
||||||
|
write_token_file("wrapper-sized-post", "tok-live", FRESH);
|
||||||
|
let inner = FakeInner::default();
|
||||||
|
let client = McpOAuthClient::new(inner.clone(), "wrapper-sized-post");
|
||||||
|
|
||||||
|
let result = rt().block_on(client.post_message_with_max_sse_event_size(
|
||||||
|
Arc::from("http://mcp.test/mcp"),
|
||||||
|
ping(),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
HashMap::new(),
|
||||||
|
4096,
|
||||||
|
));
|
||||||
|
|
||||||
|
assert!(result.is_ok());
|
||||||
|
assert_eq!(
|
||||||
|
*inner.calls.lock(),
|
||||||
|
vec![(
|
||||||
|
"post_message_with_max_sse_event_size",
|
||||||
|
Some("tok-live".to_string())
|
||||||
|
)]
|
||||||
|
);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[serial]
|
||||||
|
fn sized_get_stream_delegates_to_inner_sized_variant() {
|
||||||
|
with_temp_cache(|| {
|
||||||
|
write_token_file("wrapper-sized-get", "tok-live", FRESH);
|
||||||
|
let inner = FakeInner::default();
|
||||||
|
let client = McpOAuthClient::new(inner.clone(), "wrapper-sized-get");
|
||||||
|
|
||||||
|
let result = rt().block_on(client.get_stream_with_max_sse_event_size(
|
||||||
|
Arc::from("http://mcp.test/mcp"),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
HashMap::new(),
|
||||||
|
4096,
|
||||||
|
));
|
||||||
|
|
||||||
|
assert!(result.is_ok());
|
||||||
|
assert_eq!(
|
||||||
|
*inner.calls.lock(),
|
||||||
|
vec![(
|
||||||
|
"get_stream_with_max_sse_event_size",
|
||||||
|
Some("tok-live".to_string())
|
||||||
|
)]
|
||||||
|
);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[serial]
|
||||||
|
fn get_stream_does_not_retry_on_rejection() {
|
||||||
|
with_temp_cache(|| {
|
||||||
|
write_token_file("wrapper-get-401", "tok-live", FRESH);
|
||||||
|
let inner = FakeInner::default();
|
||||||
|
inner.reject_times.store(1, Ordering::SeqCst);
|
||||||
|
let client = McpOAuthClient::new(inner.clone(), "wrapper-get-401");
|
||||||
|
|
||||||
|
let result = rt().block_on(client.get_stream(
|
||||||
|
Arc::from("http://mcp.test/mcp"),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
HashMap::new(),
|
||||||
|
));
|
||||||
|
|
||||||
|
assert!(matches!(result, Err(StreamableHttpError::AuthRequired(_))));
|
||||||
|
assert_eq!(inner.calls.lock().len(), 1);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[serial]
|
||||||
|
fn delete_session_injects_token() {
|
||||||
|
with_temp_cache(|| {
|
||||||
|
write_token_file("wrapper-delete", "tok-live", FRESH);
|
||||||
|
let inner = FakeInner::default();
|
||||||
|
let client = McpOAuthClient::new(inner.clone(), "wrapper-delete");
|
||||||
|
|
||||||
|
let result = rt().block_on(client.delete_session(
|
||||||
|
Arc::from("http://mcp.test/mcp"),
|
||||||
|
Arc::from("session-1"),
|
||||||
|
None,
|
||||||
|
HashMap::new(),
|
||||||
|
));
|
||||||
|
|
||||||
|
assert!(result.is_ok());
|
||||||
|
assert_eq!(
|
||||||
|
*inner.calls.lock(),
|
||||||
|
vec![("delete_session", Some("tok-live".to_string()))]
|
||||||
|
);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[serial]
|
||||||
|
fn auth_required_error_contains_no_token_material() {
|
||||||
|
with_temp_cache(|| {
|
||||||
|
write_token_file("wrapper-redact", "stale-secret-token", 0);
|
||||||
|
let inner = FakeInner::default();
|
||||||
|
let client = McpOAuthClient::new(inner.clone(), "wrapper-redact");
|
||||||
|
|
||||||
|
let err = post(&client, None).unwrap_err();
|
||||||
|
|
||||||
|
let display = format!("{err}");
|
||||||
|
let debug = format!("{err:?}");
|
||||||
|
assert!(!display.contains("stale-secret-token"));
|
||||||
|
assert!(!debug.contains("stale-secret-token"));
|
||||||
|
assert!(inner.calls.lock().is_empty());
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
+314
-13
@@ -1,3 +1,4 @@
|
|||||||
|
mod auth_client;
|
||||||
pub(crate) mod manage;
|
pub(crate) mod manage;
|
||||||
pub(crate) mod oauth;
|
pub(crate) mod oauth;
|
||||||
mod sse_transport;
|
mod sse_transport;
|
||||||
@@ -9,6 +10,7 @@ use crate::vault::Vault;
|
|||||||
use crate::vault::interpolate_secrets;
|
use crate::vault::interpolate_secrets;
|
||||||
use anyhow::Error;
|
use anyhow::Error;
|
||||||
use anyhow::{Context, Result, anyhow};
|
use anyhow::{Context, Result, anyhow};
|
||||||
|
use auth_client::McpOAuthClient;
|
||||||
use futures_util::{StreamExt, TryStreamExt, stream};
|
use futures_util::{StreamExt, TryStreamExt, stream};
|
||||||
use http::{HeaderName, HeaderValue};
|
use http::{HeaderName, HeaderValue};
|
||||||
use indexmap::IndexMap;
|
use indexmap::IndexMap;
|
||||||
@@ -21,6 +23,8 @@ use rmcp::{RoleClient, ServiceExt};
|
|||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use sse_transport::LegacySseTransport;
|
use sse_transport::LegacySseTransport;
|
||||||
use std::collections::{HashMap, HashSet};
|
use std::collections::{HashMap, HashSet};
|
||||||
|
use std::fmt;
|
||||||
|
use std::fmt::Display;
|
||||||
use std::fs::OpenOptions;
|
use std::fs::OpenOptions;
|
||||||
use std::path::{Path, PathBuf};
|
use std::path::{Path, PathBuf};
|
||||||
use std::process::Stdio;
|
use std::process::Stdio;
|
||||||
@@ -326,18 +330,17 @@ impl McpRegistry {
|
|||||||
.and_then(|c| c.mcp_servers.get(&id))
|
.and_then(|c| c.mcp_servers.get(&id))
|
||||||
.with_context(|| format!("MCP server not found in config: {id}"))?;
|
.with_context(|| format!("MCP server not found in config: {id}"))?;
|
||||||
|
|
||||||
let bearer_token = if spec.is_remote() {
|
let (auth, auth_reason) = resolve_http_auth(&id, spec).await;
|
||||||
oauth::load_valid_mcp_token(&id)
|
|
||||||
} else {
|
|
||||||
None
|
|
||||||
};
|
|
||||||
|
|
||||||
let service = match spawn_mcp_server(spec, self.log_path.as_deref(), bearer_token).await {
|
let service = match spawn_mcp_server(spec, self.log_path.as_deref(), auth).await {
|
||||||
Ok(s) => s,
|
Ok(s) => s,
|
||||||
Err(e) if is_auth_required_error(&e) => {
|
Err(e) if is_auth_required_error(&e) => {
|
||||||
warn!(
|
warn!(
|
||||||
"MCP server '{id}' requires OAuth authentication. \
|
"{}",
|
||||||
Run `coyote --auth-mcp {id}` or `.mcp auth {id}` in the REPL to authenticate."
|
McpAuthRequired {
|
||||||
|
server: id,
|
||||||
|
reason: auth_reason,
|
||||||
|
}
|
||||||
);
|
);
|
||||||
return Ok(None);
|
return Ok(None);
|
||||||
}
|
}
|
||||||
@@ -412,19 +415,76 @@ impl McpRegistry {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// How a remote MCP server authenticates outgoing requests.
|
||||||
|
pub(crate) enum HttpAuth {
|
||||||
|
/// Only the static headers from the server spec; no OAuth token.
|
||||||
|
StaticOnly,
|
||||||
|
/// OAuth-managed: HTTP transports inject a fresh bearer token per request
|
||||||
|
/// via [`McpOAuthClient`] (ignoring the carried token); SSE transports
|
||||||
|
/// send the carried token as a static header.
|
||||||
|
Managed { server: String, token: String },
|
||||||
|
}
|
||||||
|
|
||||||
|
impl fmt::Debug for HttpAuth {
|
||||||
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||||
|
match self {
|
||||||
|
Self::StaticOnly => f.write_str("StaticOnly"),
|
||||||
|
Self::Managed { server, token: _ } => f
|
||||||
|
.debug_struct("Managed")
|
||||||
|
.field("server", server)
|
||||||
|
.field("token", &"<redacted>")
|
||||||
|
.finish(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl HttpAuth {
|
||||||
|
pub(crate) fn from_token_status(status: &oauth::McpTokenStatus, server: &str) -> Self {
|
||||||
|
match status {
|
||||||
|
oauth::McpTokenStatus::Token(token) => Self::Managed {
|
||||||
|
server: server.to_string(),
|
||||||
|
token: token.clone(),
|
||||||
|
},
|
||||||
|
oauth::McpTokenStatus::NotAuthenticated | oauth::McpTokenStatus::RefreshFailed => {
|
||||||
|
Self::StaticOnly
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn resolve_http_auth(name: &str, spec: &McpServer) -> (HttpAuth, McpAuthReason) {
|
||||||
|
let token_status = if spec.is_remote() {
|
||||||
|
oauth::load_or_refresh_mcp_token(name).await
|
||||||
|
} else {
|
||||||
|
oauth::McpTokenStatus::NotAuthenticated
|
||||||
|
};
|
||||||
|
(
|
||||||
|
HttpAuth::from_token_status(&token_status, name),
|
||||||
|
McpAuthReason::from_token_status(&token_status),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
pub(crate) async fn spawn_mcp_server(
|
pub(crate) async fn spawn_mcp_server(
|
||||||
spec: &McpServer,
|
spec: &McpServer,
|
||||||
log_path: Option<&Path>,
|
log_path: Option<&Path>,
|
||||||
bearer_token: Option<String>,
|
auth: HttpAuth,
|
||||||
) -> Result<Arc<ConnectedServer>> {
|
) -> Result<Arc<ConnectedServer>> {
|
||||||
match spec.transport_type {
|
match spec.transport_type {
|
||||||
McpTransportType::Http => {
|
McpTransportType::Http => {
|
||||||
let url = spec.url.as_deref().expect("validated: http spec has url");
|
let url = spec.url.as_deref().expect("validated: http spec has url");
|
||||||
let headers = merge_bearer_token(spec.headers.as_ref(), bearer_token);
|
match auth {
|
||||||
spawn_http_mcp_server(url, headers.as_ref()).await
|
HttpAuth::Managed { server, token: _ } => {
|
||||||
|
spawn_oauth_http_mcp_server(url, &server, spec.headers.as_ref()).await
|
||||||
|
}
|
||||||
|
HttpAuth::StaticOnly => spawn_http_mcp_server(url, spec.headers.as_ref()).await,
|
||||||
|
}
|
||||||
}
|
}
|
||||||
McpTransportType::Sse => {
|
McpTransportType::Sse => {
|
||||||
let url = spec.url.as_deref().expect("validated: sse spec has url");
|
let url = spec.url.as_deref().expect("validated: sse spec has url");
|
||||||
|
let bearer_token = match auth {
|
||||||
|
HttpAuth::Managed { server: _, token } => Some(token),
|
||||||
|
HttpAuth::StaticOnly => None,
|
||||||
|
};
|
||||||
let headers = merge_bearer_token(spec.headers.as_ref(), bearer_token);
|
let headers = merge_bearer_token(spec.headers.as_ref(), bearer_token);
|
||||||
spawn_sse_mcp_server(url, headers.as_ref()).await
|
spawn_sse_mcp_server(url, headers.as_ref()).await
|
||||||
}
|
}
|
||||||
@@ -452,15 +512,66 @@ fn merge_bearer_token(
|
|||||||
}
|
}
|
||||||
(Some(h), Some(token)) => {
|
(Some(h), Some(token)) => {
|
||||||
let mut m = h.clone();
|
let mut m = h.clone();
|
||||||
|
m.retain(|k, _| !k.eq_ignore_ascii_case("authorization"));
|
||||||
m.insert("Authorization".to_string(), format!("Bearer {token}"));
|
m.insert("Authorization".to_string(), format!("Bearer {token}"));
|
||||||
Some(m)
|
Some(m)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||||
|
pub(crate) enum McpAuthReason {
|
||||||
|
NotAuthenticated,
|
||||||
|
RefreshFailed,
|
||||||
|
TokenRejected,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl McpAuthReason {
|
||||||
|
pub(crate) fn from_token_status(status: &oauth::McpTokenStatus) -> Self {
|
||||||
|
match status {
|
||||||
|
oauth::McpTokenStatus::Token(_) => Self::TokenRejected,
|
||||||
|
oauth::McpTokenStatus::NotAuthenticated => Self::NotAuthenticated,
|
||||||
|
oauth::McpTokenStatus::RefreshFailed => Self::RefreshFailed,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub(crate) struct McpAuthRequired {
|
||||||
|
pub server: String,
|
||||||
|
pub reason: McpAuthReason,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Display for McpAuthRequired {
|
||||||
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||||
|
let server = &self.server;
|
||||||
|
match self.reason {
|
||||||
|
McpAuthReason::NotAuthenticated => write!(
|
||||||
|
f,
|
||||||
|
"MCP server '{server}' requires OAuth authentication and was not started \
|
||||||
|
(no stored credentials). Run `.mcp auth {server}` (or `coyote --auth-mcp \
|
||||||
|
{server}`) to authenticate and attach it."
|
||||||
|
),
|
||||||
|
McpAuthReason::RefreshFailed => write!(
|
||||||
|
f,
|
||||||
|
"MCP server '{server}' was not started: stored OAuth token has expired and \
|
||||||
|
automatic refresh failed. Run `.mcp auth {server}` (or `coyote --auth-mcp \
|
||||||
|
{server}`) to re-authenticate and attach it."
|
||||||
|
),
|
||||||
|
McpAuthReason::TokenRejected => write!(
|
||||||
|
f,
|
||||||
|
"MCP server '{server}' was not started: the server rejected the stored OAuth \
|
||||||
|
token. Run `.mcp auth {server}` (or `coyote --auth-mcp {server}`) to \
|
||||||
|
re-authenticate and attach it."
|
||||||
|
),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub(crate) fn is_auth_required_error(e: &Error) -> bool {
|
pub(crate) fn is_auth_required_error(e: &Error) -> bool {
|
||||||
e.chain()
|
e.downcast_ref::<McpAuthRequired>().is_some()
|
||||||
.any(|cause| cause.to_string().contains("Auth required"))
|
|| e.chain()
|
||||||
|
.any(|cause| cause.to_string().contains("Auth required"))
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn spawn_http_mcp_server(
|
async fn spawn_http_mcp_server(
|
||||||
@@ -493,6 +604,66 @@ async fn spawn_http_mcp_server(
|
|||||||
Ok(service)
|
Ok(service)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Builds the custom-header map for an OAuth-managed HTTP transport, dropping
|
||||||
|
/// any static `Authorization` entry case-insensitively: [`McpOAuthClient`]
|
||||||
|
/// owns that header, and a stale configured value must not collide with the
|
||||||
|
/// per-request token.
|
||||||
|
fn oauth_custom_headers(
|
||||||
|
headers: Option<&IndexMap<String, String>>,
|
||||||
|
) -> Result<HashMap<HeaderName, HeaderValue>> {
|
||||||
|
let mut custom = HashMap::new();
|
||||||
|
let Some(hdrs) = headers else {
|
||||||
|
return Ok(custom);
|
||||||
|
};
|
||||||
|
|
||||||
|
for (k, v) in hdrs {
|
||||||
|
if k.eq_ignore_ascii_case("authorization") {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let name = k
|
||||||
|
.parse::<HeaderName>()
|
||||||
|
.with_context(|| format!("Invalid header name: {k}"))?;
|
||||||
|
let value = v
|
||||||
|
.parse::<HeaderValue>()
|
||||||
|
.with_context(|| format!("Invalid header value for {k}"))?;
|
||||||
|
custom.insert(name, value);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(custom)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn spawn_oauth_http_mcp_server(
|
||||||
|
url: &str,
|
||||||
|
server: &str,
|
||||||
|
headers: Option<&IndexMap<String, String>>,
|
||||||
|
) -> Result<Arc<ConnectedServer>> {
|
||||||
|
// Mirror rmcp's default_http_client, which `with_client` bypasses:
|
||||||
|
// idle pooling off avoids a documented TCP delayed-ACK stall, and
|
||||||
|
// redirects off keeps custom headers from being replayed to a redirect
|
||||||
|
// target.
|
||||||
|
let inner = reqwest::Client::builder()
|
||||||
|
.pool_max_idle_per_host(0)
|
||||||
|
.redirect(reqwest::redirect::Policy::none())
|
||||||
|
.build()
|
||||||
|
.context("Failed to build HTTP client for OAuth-managed MCP transport")?;
|
||||||
|
let client = McpOAuthClient::new(inner, server);
|
||||||
|
// `auth_header` stays None so the wrapper injects a fresh token per
|
||||||
|
// request; `reinit_on_expired_session` defaults to true in rmcp 3.1.2
|
||||||
|
// but is pinned explicitly because transparent session re-init is
|
||||||
|
// load-bearing for long-lived sessions.
|
||||||
|
let config = StreamableHttpClientTransportConfig::with_uri(url)
|
||||||
|
.custom_headers(oauth_custom_headers(headers)?)
|
||||||
|
.reinit_on_expired_session(true);
|
||||||
|
let transport = StreamableHttpClientTransport::with_client(client, config);
|
||||||
|
let service = Arc::new(
|
||||||
|
().serve(transport)
|
||||||
|
.await
|
||||||
|
.with_context(|| format!("Failed to connect to HTTP MCP server: {url}"))?,
|
||||||
|
);
|
||||||
|
|
||||||
|
Ok(service)
|
||||||
|
}
|
||||||
|
|
||||||
async fn spawn_sse_mcp_server(
|
async fn spawn_sse_mcp_server(
|
||||||
url: &str,
|
url: &str,
|
||||||
headers: Option<&IndexMap<String, String>>,
|
headers: Option<&IndexMap<String, String>>,
|
||||||
@@ -1051,6 +1222,92 @@ mod tests {
|
|||||||
assert_eq!(result["X-Custom"], "keep");
|
assert_eq!(result["X-Custom"], "keep");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn merge_bearer_token_replaces_authorization_case_insensitively() {
|
||||||
|
let mut h = IndexMap::new();
|
||||||
|
h.insert("authorization".to_string(), "Bearer stale-1".to_string());
|
||||||
|
h.insert("AUTHORIZATION".to_string(), "Bearer stale-2".to_string());
|
||||||
|
h.insert("X-Custom".to_string(), "keep".to_string());
|
||||||
|
|
||||||
|
let result = merge_bearer_token(Some(&h), Some("newtoken".to_string())).unwrap();
|
||||||
|
|
||||||
|
assert_eq!(result.len(), 2);
|
||||||
|
assert_eq!(result["Authorization"], "Bearer newtoken");
|
||||||
|
assert_eq!(result["X-Custom"], "keep");
|
||||||
|
assert!(!result.contains_key("authorization"));
|
||||||
|
assert!(!result.contains_key("AUTHORIZATION"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn http_auth_from_token_status_maps_token_to_managed() {
|
||||||
|
assert!(matches!(
|
||||||
|
HttpAuth::from_token_status(&oauth::McpTokenStatus::Token("tok".into()), "srv"),
|
||||||
|
HttpAuth::Managed { server, token } if server == "srv" && token == "tok"
|
||||||
|
));
|
||||||
|
assert!(matches!(
|
||||||
|
HttpAuth::from_token_status(&oauth::McpTokenStatus::NotAuthenticated, "srv"),
|
||||||
|
HttpAuth::StaticOnly
|
||||||
|
));
|
||||||
|
assert!(matches!(
|
||||||
|
HttpAuth::from_token_status(&oauth::McpTokenStatus::RefreshFailed, "srv"),
|
||||||
|
HttpAuth::StaticOnly
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn http_auth_debug_redacts_token() {
|
||||||
|
let auth = HttpAuth::Managed {
|
||||||
|
server: "srv".into(),
|
||||||
|
token: "live-secret".into(),
|
||||||
|
};
|
||||||
|
|
||||||
|
let debug = format!("{auth:?}");
|
||||||
|
|
||||||
|
assert!(debug.contains("srv"));
|
||||||
|
assert!(debug.contains("<redacted>"));
|
||||||
|
assert!(!debug.contains("live-secret"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn oauth_custom_headers_strips_authorization_case_insensitively() {
|
||||||
|
let mut h = IndexMap::new();
|
||||||
|
h.insert("Authorization".to_string(), "Bearer stale-1".to_string());
|
||||||
|
h.insert("authorization".to_string(), "Bearer stale-2".to_string());
|
||||||
|
h.insert("AUTHORIZATION".to_string(), "Bearer stale-3".to_string());
|
||||||
|
h.insert("X-Custom".to_string(), "keep".to_string());
|
||||||
|
|
||||||
|
let custom = oauth_custom_headers(Some(&h)).unwrap();
|
||||||
|
|
||||||
|
assert_eq!(custom.len(), 1);
|
||||||
|
assert_eq!(custom[&HeaderName::from_static("x-custom")], "keep");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn oauth_custom_headers_none_is_empty() {
|
||||||
|
assert!(oauth_custom_headers(None).unwrap().is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn oauth_custom_headers_rejects_invalid_header_name() {
|
||||||
|
let mut h = IndexMap::new();
|
||||||
|
h.insert("bad header".to_string(), "v".to_string());
|
||||||
|
|
||||||
|
assert!(oauth_custom_headers(Some(&h)).is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn oauth_custom_headers_keeps_non_authorization_headers() {
|
||||||
|
let mut h = IndexMap::new();
|
||||||
|
h.insert("X-Api-Key".to_string(), "k".to_string());
|
||||||
|
h.insert("X-Trace".to_string(), "t".to_string());
|
||||||
|
|
||||||
|
let custom = oauth_custom_headers(Some(&h)).unwrap();
|
||||||
|
|
||||||
|
assert_eq!(custom.len(), 2);
|
||||||
|
assert_eq!(custom[&HeaderName::from_static("x-api-key")], "k");
|
||||||
|
assert_eq!(custom[&HeaderName::from_static("x-trace")], "t");
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn is_auth_required_error_matches_rmcp_message() {
|
fn is_auth_required_error_matches_rmcp_message() {
|
||||||
let e = anyhow!("Auth required, when send initialize request");
|
let e = anyhow!("Auth required, when send initialize request");
|
||||||
@@ -1074,4 +1331,48 @@ mod tests {
|
|||||||
|
|
||||||
assert!(is_auth_required_error(&e));
|
assert!(is_auth_required_error(&e));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn auth_reason_maps_token_status() {
|
||||||
|
assert_eq!(
|
||||||
|
McpAuthReason::from_token_status(&oauth::McpTokenStatus::Token("tok".into())),
|
||||||
|
McpAuthReason::TokenRejected
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
McpAuthReason::from_token_status(&oauth::McpTokenStatus::NotAuthenticated),
|
||||||
|
McpAuthReason::NotAuthenticated
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
McpAuthReason::from_token_status(&oauth::McpTokenStatus::RefreshFailed),
|
||||||
|
McpAuthReason::RefreshFailed
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn mcp_auth_required_context_downcasts_with_reason() {
|
||||||
|
let e = anyhow!("Auth required, when send initialize request").context(McpAuthRequired {
|
||||||
|
server: "github".into(),
|
||||||
|
reason: McpAuthReason::RefreshFailed,
|
||||||
|
});
|
||||||
|
|
||||||
|
assert!(is_auth_required_error(&e));
|
||||||
|
let ctx = e.downcast_ref::<McpAuthRequired>().unwrap();
|
||||||
|
assert_eq!(ctx.server, "github");
|
||||||
|
assert_eq!(ctx.reason, McpAuthReason::RefreshFailed);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn mcp_auth_required_display_is_reason_specific() {
|
||||||
|
let msg = |reason| {
|
||||||
|
McpAuthRequired {
|
||||||
|
server: "github".into(),
|
||||||
|
reason,
|
||||||
|
}
|
||||||
|
.to_string()
|
||||||
|
};
|
||||||
|
|
||||||
|
assert!(msg(McpAuthReason::NotAuthenticated).contains("no stored credentials"));
|
||||||
|
assert!(msg(McpAuthReason::RefreshFailed).contains("expired and automatic refresh failed"));
|
||||||
|
assert!(msg(McpAuthReason::TokenRejected).contains("rejected the stored OAuth token"));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+532
-33
@@ -1,15 +1,26 @@
|
|||||||
use crate::client::oauth::{OAuthProvider, TokenRequestFormat, load_oauth_tokens, run_oauth_flow};
|
use crate::client::oauth::{
|
||||||
|
OAuthProvider, OAuthTokens, TokenRequestFormat, load_oauth_tokens, refresh_oauth_token,
|
||||||
|
run_oauth_flow, token_response_keys,
|
||||||
|
};
|
||||||
use crate::config::paths;
|
use crate::config::paths;
|
||||||
use anyhow::{Context, Result, anyhow};
|
use anyhow::{Context, Result, anyhow};
|
||||||
use chrono::Utc;
|
use chrono::Utc;
|
||||||
use inquire::Text;
|
use inquire::Text;
|
||||||
use log::warn;
|
use log::{debug, warn};
|
||||||
use reqwest::Client;
|
use reqwest::Client;
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::fmt;
|
||||||
use std::fs;
|
use std::fs;
|
||||||
use std::net::TcpListener;
|
use std::net::TcpListener;
|
||||||
|
use std::sync::{Arc, OnceLock};
|
||||||
|
use std::time::{Duration, Instant};
|
||||||
|
use tokio::sync;
|
||||||
use url::Url;
|
use url::Url;
|
||||||
|
|
||||||
|
const REFRESH_HTTP_TIMEOUT: Duration = Duration::from_secs(10);
|
||||||
|
const REFRESH_FAILURE_BACKOFF: Duration = Duration::from_secs(60);
|
||||||
|
|
||||||
#[derive(Debug, Deserialize)]
|
#[derive(Debug, Deserialize)]
|
||||||
struct ProtectedResourceMetadata {
|
struct ProtectedResourceMetadata {
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
@@ -34,6 +45,10 @@ struct McpRegistration {
|
|||||||
client_id: String,
|
client_id: String,
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
redirect_uri: Option<String>,
|
redirect_uri: Option<String>,
|
||||||
|
#[serde(default)]
|
||||||
|
token_url: Option<String>,
|
||||||
|
#[serde(default)]
|
||||||
|
resource: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
struct DiscoveredOAuth {
|
struct DiscoveredOAuth {
|
||||||
@@ -124,8 +139,19 @@ pub async fn run_mcp_oauth_flow(
|
|||||||
None
|
None
|
||||||
};
|
};
|
||||||
|
|
||||||
let (client_id, redirect_uri) = if let Some(reused) = cached_reuse {
|
let (client_id, redirect_uri) = if let Some((client_id, redirect_uri)) = cached_reuse {
|
||||||
reused
|
// Re-save so registrations cached before token_url/resource were
|
||||||
|
// persisted gain them, enabling token refresh next time.
|
||||||
|
if let Err(e) = save_registration(
|
||||||
|
server_name,
|
||||||
|
&client_id,
|
||||||
|
&redirect_uri,
|
||||||
|
&metadata.token_endpoint,
|
||||||
|
&resource,
|
||||||
|
) {
|
||||||
|
debug!("Failed to update cached MCP registration for '{server_name}': {e}");
|
||||||
|
}
|
||||||
|
(client_id, redirect_uri)
|
||||||
} else {
|
} else {
|
||||||
let bind_addr = format!("127.0.0.1:{}", callback_port.unwrap_or(0));
|
let bind_addr = format!("127.0.0.1:{}", callback_port.unwrap_or(0));
|
||||||
let listener = TcpListener::bind(&bind_addr)?;
|
let listener = TcpListener::bind(&bind_addr)?;
|
||||||
@@ -137,10 +163,7 @@ pub async fn run_mcp_oauth_flow(
|
|||||||
id.to_string()
|
id.to_string()
|
||||||
} else if let Some(reg_endpoint) = &metadata.registration_endpoint {
|
} else if let Some(reg_endpoint) = &metadata.registration_endpoint {
|
||||||
match register_client(reg_endpoint, &redirect_uri).await {
|
match register_client(reg_endpoint, &redirect_uri).await {
|
||||||
Ok(id) => {
|
Ok(id) => id,
|
||||||
let _ = save_registration(server_name, &id, &redirect_uri);
|
|
||||||
id
|
|
||||||
}
|
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
warn!("Dynamic client registration failed: {e}. Falling back to manual entry.");
|
warn!("Dynamic client registration failed: {e}. Falling back to manual entry.");
|
||||||
Text::new("Enter the OAuth client ID for this MCP server:")
|
Text::new("Enter the OAuth client ID for this MCP server:")
|
||||||
@@ -153,6 +176,18 @@ pub async fn run_mcp_oauth_flow(
|
|||||||
.prompt()
|
.prompt()
|
||||||
.context("Failed to read client ID")?
|
.context("Failed to read client ID")?
|
||||||
};
|
};
|
||||||
|
// Persist regardless of how the client_id was obtained (DCR, config,
|
||||||
|
// or manual entry) so refresh_mcp_token can run the refresh_token
|
||||||
|
// grant later without interactive re-auth.
|
||||||
|
if let Err(e) = save_registration(
|
||||||
|
server_name,
|
||||||
|
&client_id,
|
||||||
|
&redirect_uri,
|
||||||
|
&metadata.token_endpoint,
|
||||||
|
&resource,
|
||||||
|
) {
|
||||||
|
debug!("Failed to cache MCP registration for '{server_name}': {e}");
|
||||||
|
}
|
||||||
(client_id, redirect_uri)
|
(client_id, redirect_uri)
|
||||||
};
|
};
|
||||||
|
|
||||||
@@ -168,12 +203,164 @@ pub async fn run_mcp_oauth_flow(
|
|||||||
run_oauth_flow(&provider, &mcp_token_key(server_name)).await
|
run_oauth_flow(&provider, &mcp_token_key(server_name)).await
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn load_valid_mcp_token(server_name: &str) -> Option<String> {
|
#[derive(PartialEq, Eq)]
|
||||||
let tokens = load_oauth_tokens(&mcp_token_key(server_name))?;
|
pub enum McpTokenStatus {
|
||||||
if Utc::now().timestamp() < tokens.expires_at {
|
Token(String),
|
||||||
Some(tokens.access_token)
|
NotAuthenticated,
|
||||||
} else {
|
RefreshFailed,
|
||||||
None
|
}
|
||||||
|
|
||||||
|
impl fmt::Debug for McpTokenStatus {
|
||||||
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||||
|
match self {
|
||||||
|
Self::Token(_) => f.write_str("Token(<redacted>)"),
|
||||||
|
Self::NotAuthenticated => f.write_str("NotAuthenticated"),
|
||||||
|
Self::RefreshFailed => f.write_str("RefreshFailed"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl McpTokenStatus {
|
||||||
|
pub fn into_token(self) -> Option<String> {
|
||||||
|
match self {
|
||||||
|
Self::Token(token) => Some(token),
|
||||||
|
Self::NotAuthenticated | Self::RefreshFailed => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn load_or_refresh_mcp_token(server_name: &str) -> McpTokenStatus {
|
||||||
|
load_or_refresh_inner(server_name, None).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Re-acquires a token after the server rejected the current one mid-session
|
||||||
|
/// (HTTP 401). The rejection proves the stored token is bad regardless of its
|
||||||
|
/// expiry timestamp, so the unexpired fast-paths only short-circuit when the
|
||||||
|
/// stored token DIFFERS from `rejected_token` (a concurrent caller genuinely
|
||||||
|
/// refreshed while we waited); an unexpired copy of the rejected token is
|
||||||
|
/// refreshed anyway. The failure backoff and per-server single-flight lock
|
||||||
|
/// still apply.
|
||||||
|
pub async fn force_refresh_mcp_token(server_name: &str, rejected_token: &str) -> Option<String> {
|
||||||
|
load_or_refresh_inner(server_name, Some(rejected_token))
|
||||||
|
.await
|
||||||
|
.into_token()
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn load_or_refresh_inner(server_name: &str, rejected_token: Option<&str>) -> McpTokenStatus {
|
||||||
|
let key = mcp_token_key(server_name);
|
||||||
|
let Some(tokens) = load_oauth_tokens(&key) else {
|
||||||
|
return McpTokenStatus::NotAuthenticated;
|
||||||
|
};
|
||||||
|
if rejected_token.is_none() && Utc::now().timestamp() < tokens.expires_at {
|
||||||
|
return McpTokenStatus::Token(tokens.access_token);
|
||||||
|
}
|
||||||
|
|
||||||
|
if in_refresh_failure_backoff(server_name) {
|
||||||
|
debug!("Skipping token refresh for MCP server '{server_name}': recent attempt failed");
|
||||||
|
return McpTokenStatus::RefreshFailed;
|
||||||
|
}
|
||||||
|
|
||||||
|
let lock = refresh_lock(server_name);
|
||||||
|
let _guard = lock.lock().await;
|
||||||
|
|
||||||
|
// A concurrent caller may have refreshed while we waited for the lock. An
|
||||||
|
// unexpired token is only trusted if it differs from the rejected one:
|
||||||
|
// the server already proved that exact token bad.
|
||||||
|
let Some(tokens) = load_oauth_tokens(&key) else {
|
||||||
|
return McpTokenStatus::NotAuthenticated;
|
||||||
|
};
|
||||||
|
if Utc::now().timestamp() < tokens.expires_at
|
||||||
|
&& rejected_token.is_none_or(|rejected| rejected != tokens.access_token)
|
||||||
|
{
|
||||||
|
return McpTokenStatus::Token(tokens.access_token);
|
||||||
|
}
|
||||||
|
|
||||||
|
if in_refresh_failure_backoff(server_name) {
|
||||||
|
debug!("Skipping token refresh for MCP server '{server_name}': recent attempt failed");
|
||||||
|
return McpTokenStatus::RefreshFailed;
|
||||||
|
}
|
||||||
|
|
||||||
|
match refresh_mcp_token(server_name, &key, &tokens).await {
|
||||||
|
Ok(access_token) => McpTokenStatus::Token(access_token),
|
||||||
|
Err(e) => {
|
||||||
|
note_refresh_failure(server_name);
|
||||||
|
warn!(
|
||||||
|
"Failed to refresh OAuth token for MCP server '{server_name}'. \
|
||||||
|
Run `.mcp auth {server_name}` to re-authenticate."
|
||||||
|
);
|
||||||
|
debug!(
|
||||||
|
"Token refresh error for MCP server '{server_name}': {}",
|
||||||
|
redact_refresh_error(&e)
|
||||||
|
);
|
||||||
|
McpTokenStatus::RefreshFailed
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn refresh_mcp_token(server_name: &str, key: &str, tokens: &OAuthTokens) -> Result<String> {
|
||||||
|
if tokens.refresh_token.is_none() {
|
||||||
|
return Err(anyhow!("no refresh token stored"));
|
||||||
|
}
|
||||||
|
|
||||||
|
let reg =
|
||||||
|
load_registration(server_name).ok_or_else(|| anyhow!("no cached client registration"))?;
|
||||||
|
let token_url = reg.token_url.ok_or_else(|| {
|
||||||
|
anyhow!("cached registration has no token URL (saved by an older version)")
|
||||||
|
})?;
|
||||||
|
let resource = reg.resource.ok_or_else(|| {
|
||||||
|
anyhow!("cached registration has no resource (saved by an older version)")
|
||||||
|
})?;
|
||||||
|
|
||||||
|
let provider = McpOAuthProvider {
|
||||||
|
client_id: reg.client_id,
|
||||||
|
authorize_url: String::new(),
|
||||||
|
token_url,
|
||||||
|
scopes: String::new(),
|
||||||
|
fixed_redirect: String::new(),
|
||||||
|
resource,
|
||||||
|
};
|
||||||
|
|
||||||
|
let client = Client::builder().timeout(REFRESH_HTTP_TIMEOUT).build()?;
|
||||||
|
let refreshed = refresh_oauth_token(&client, &provider, key, tokens).await?;
|
||||||
|
Ok(refreshed.access_token)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn refresh_lock(server_name: &str) -> Arc<sync::Mutex<()>> {
|
||||||
|
static LOCKS: OnceLock<parking_lot::Mutex<HashMap<String, Arc<sync::Mutex<()>>>>> =
|
||||||
|
OnceLock::new();
|
||||||
|
LOCKS
|
||||||
|
.get_or_init(Default::default)
|
||||||
|
.lock()
|
||||||
|
.entry(server_name.to_string())
|
||||||
|
.or_default()
|
||||||
|
.clone()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn refresh_failures() -> &'static parking_lot::Mutex<HashMap<String, Instant>> {
|
||||||
|
static FAILURES: OnceLock<parking_lot::Mutex<HashMap<String, Instant>>> = OnceLock::new();
|
||||||
|
FAILURES.get_or_init(Default::default)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn note_refresh_failure(server_name: &str) {
|
||||||
|
refresh_failures()
|
||||||
|
.lock()
|
||||||
|
.insert(server_name.to_string(), Instant::now());
|
||||||
|
}
|
||||||
|
|
||||||
|
fn in_refresh_failure_backoff(server_name: &str) -> bool {
|
||||||
|
refresh_failures()
|
||||||
|
.lock()
|
||||||
|
.get(server_name)
|
||||||
|
.is_some_and(|failed_at| failed_at.elapsed() < REFRESH_FAILURE_BACKOFF)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Refresh errors may embed the token endpoint's JSON response, which can
|
||||||
|
/// contain live tokens; strip everything from the first `{` before logging.
|
||||||
|
fn redact_refresh_error(e: &anyhow::Error) -> String {
|
||||||
|
let msg = e.to_string();
|
||||||
|
match msg.find('{') {
|
||||||
|
Some(idx) => format!("{}<response body redacted>", &msg[..idx]),
|
||||||
|
None => msg,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -187,7 +374,13 @@ fn load_registration(server_name: &str) -> Option<McpRegistration> {
|
|||||||
serde_json::from_str(&content).ok()
|
serde_json::from_str(&content).ok()
|
||||||
}
|
}
|
||||||
|
|
||||||
fn save_registration(server_name: &str, client_id: &str, redirect_uri: &str) -> Result<()> {
|
fn save_registration(
|
||||||
|
server_name: &str,
|
||||||
|
client_id: &str,
|
||||||
|
redirect_uri: &str,
|
||||||
|
token_url: &str,
|
||||||
|
resource: &str,
|
||||||
|
) -> Result<()> {
|
||||||
let dir = paths::oauth_tokens_dir();
|
let dir = paths::oauth_tokens_dir();
|
||||||
fs::create_dir_all(&dir)?;
|
fs::create_dir_all(&dir)?;
|
||||||
|
|
||||||
@@ -195,6 +388,8 @@ fn save_registration(server_name: &str, client_id: &str, redirect_uri: &str) ->
|
|||||||
let reg = McpRegistration {
|
let reg = McpRegistration {
|
||||||
client_id: client_id.to_string(),
|
client_id: client_id.to_string(),
|
||||||
redirect_uri: Some(redirect_uri.to_string()),
|
redirect_uri: Some(redirect_uri.to_string()),
|
||||||
|
token_url: Some(token_url.to_string()),
|
||||||
|
resource: Some(resource.to_string()),
|
||||||
};
|
};
|
||||||
|
|
||||||
fs::write(path, serde_json::to_string_pretty(®)?)?;
|
fs::write(path, serde_json::to_string_pretty(®)?)?;
|
||||||
@@ -244,7 +439,12 @@ async fn register_client(endpoint: &str, redirect_uri: &str) -> Result<String> {
|
|||||||
|
|
||||||
response["client_id"]
|
response["client_id"]
|
||||||
.as_str()
|
.as_str()
|
||||||
.ok_or_else(|| anyhow!("Missing client_id in registration response: {response}"))
|
.ok_or_else(|| {
|
||||||
|
anyhow!(
|
||||||
|
"Missing client_id in registration response (keys: {})",
|
||||||
|
token_response_keys(&response)
|
||||||
|
)
|
||||||
|
})
|
||||||
.map(|s| s.to_string())
|
.map(|s| s.to_string())
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -421,16 +621,34 @@ fn extract_base_url(url: &str) -> Result<String> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
pub(crate) mod test_support {
|
||||||
use super::*;
|
|
||||||
use crate::utils::get_env_name;
|
use crate::utils::get_env_name;
|
||||||
use serial_test::serial;
|
|
||||||
use std::{
|
use std::{
|
||||||
env, fs,
|
env,
|
||||||
|
ffi::OsString,
|
||||||
|
fs,
|
||||||
|
path::PathBuf,
|
||||||
time::{self, SystemTime},
|
time::{self, SystemTime},
|
||||||
};
|
};
|
||||||
|
|
||||||
fn with_temp_cache<F: FnOnce()>(f: F) {
|
pub(crate) fn with_temp_cache<F: FnOnce()>(f: F) {
|
||||||
|
struct Restore {
|
||||||
|
key: String,
|
||||||
|
prev: Option<OsString>,
|
||||||
|
root: PathBuf,
|
||||||
|
}
|
||||||
|
impl Drop for Restore {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
unsafe {
|
||||||
|
match self.prev.take() {
|
||||||
|
Some(v) => env::set_var(&self.key, v),
|
||||||
|
None => env::remove_var(&self.key),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let _ = fs::remove_dir_all(&self.root);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
let unique = SystemTime::now()
|
let unique = SystemTime::now()
|
||||||
.duration_since(time::UNIX_EPOCH)
|
.duration_since(time::UNIX_EPOCH)
|
||||||
.unwrap()
|
.unwrap()
|
||||||
@@ -442,15 +660,21 @@ mod tests {
|
|||||||
unsafe {
|
unsafe {
|
||||||
env::set_var(&env_key, &root);
|
env::set_var(&env_key, &root);
|
||||||
}
|
}
|
||||||
|
let _restore = Restore {
|
||||||
|
key: env_key,
|
||||||
|
prev,
|
||||||
|
root,
|
||||||
|
};
|
||||||
f();
|
f();
|
||||||
unsafe {
|
|
||||||
match prev {
|
|
||||||
Some(v) => env::set_var(&env_key, v),
|
|
||||||
None => env::remove_var(&env_key),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
let _ = fs::remove_dir_all(&root);
|
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::test_support::with_temp_cache;
|
||||||
|
use super::*;
|
||||||
|
use serial_test::serial;
|
||||||
|
use std::fs;
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn extract_base_url_strips_path_and_query() {
|
fn extract_base_url_strips_path_and_query() {
|
||||||
@@ -685,12 +909,19 @@ mod tests {
|
|||||||
"notion",
|
"notion",
|
||||||
"client-xyz-123",
|
"client-xyz-123",
|
||||||
"http://127.0.0.1:49152/callback",
|
"http://127.0.0.1:49152/callback",
|
||||||
|
"https://as.example/token",
|
||||||
|
"https://mcp.example/mcp",
|
||||||
)
|
)
|
||||||
.unwrap();
|
.unwrap();
|
||||||
|
|
||||||
let loaded = load_registration("notion");
|
let loaded = load_registration("notion").unwrap();
|
||||||
|
|
||||||
assert_eq!(loaded.unwrap().client_id, "client-xyz-123");
|
assert_eq!(loaded.client_id, "client-xyz-123");
|
||||||
|
assert_eq!(
|
||||||
|
loaded.token_url.as_deref(),
|
||||||
|
Some("https://as.example/token")
|
||||||
|
);
|
||||||
|
assert_eq!(loaded.resource.as_deref(), Some("https://mcp.example/mcp"));
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -708,8 +939,22 @@ mod tests {
|
|||||||
#[serial]
|
#[serial]
|
||||||
fn registration_second_save_overwrites_first() {
|
fn registration_second_save_overwrites_first() {
|
||||||
with_temp_cache(|| {
|
with_temp_cache(|| {
|
||||||
save_registration("github", "first-id", "http://127.0.0.1:49152/callback").unwrap();
|
save_registration(
|
||||||
save_registration("github", "second-id", "http://127.0.0.1:49153/callback").unwrap();
|
"github",
|
||||||
|
"first-id",
|
||||||
|
"http://127.0.0.1:49152/callback",
|
||||||
|
"https://as.example/token",
|
||||||
|
"https://mcp.example/mcp",
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
save_registration(
|
||||||
|
"github",
|
||||||
|
"second-id",
|
||||||
|
"http://127.0.0.1:49153/callback",
|
||||||
|
"https://as.example/token",
|
||||||
|
"https://mcp.example/mcp",
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
let loaded = load_registration("github").unwrap();
|
let loaded = load_registration("github").unwrap();
|
||||||
|
|
||||||
@@ -737,6 +982,8 @@ mod tests {
|
|||||||
|
|
||||||
assert_eq!(loaded.client_id, "legacy-id");
|
assert_eq!(loaded.client_id, "legacy-id");
|
||||||
assert_eq!(loaded.redirect_uri, None);
|
assert_eq!(loaded.redirect_uri, None);
|
||||||
|
assert_eq!(loaded.token_url, None);
|
||||||
|
assert_eq!(loaded.resource, None);
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -744,7 +991,14 @@ mod tests {
|
|||||||
#[serial]
|
#[serial]
|
||||||
fn save_registration_persists_redirect_uri() {
|
fn save_registration_persists_redirect_uri() {
|
||||||
with_temp_cache(|| {
|
with_temp_cache(|| {
|
||||||
save_registration("aws", "client-abc", "http://127.0.0.1:49152/callback").unwrap();
|
save_registration(
|
||||||
|
"aws",
|
||||||
|
"client-abc",
|
||||||
|
"http://127.0.0.1:49152/callback",
|
||||||
|
"https://as.example/token",
|
||||||
|
"https://mcp.example/mcp",
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
let loaded = load_registration("aws").unwrap();
|
let loaded = load_registration("aws").unwrap();
|
||||||
|
|
||||||
@@ -756,6 +1010,251 @@ mod tests {
|
|||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn mcp_registration_deserializes_without_new_fields_and_roundtrips() {
|
||||||
|
let old: McpRegistration = serde_json::from_str(r#"{"client_id":"legacy-id"}"#).unwrap();
|
||||||
|
|
||||||
|
assert_eq!(old.client_id, "legacy-id");
|
||||||
|
assert_eq!(old.token_url, None);
|
||||||
|
assert_eq!(old.resource, None);
|
||||||
|
|
||||||
|
let full = McpRegistration {
|
||||||
|
client_id: "client-abc".into(),
|
||||||
|
redirect_uri: Some("http://127.0.0.1:49152/callback".into()),
|
||||||
|
token_url: Some("https://as.example/token".into()),
|
||||||
|
resource: Some("https://mcp.example/mcp".into()),
|
||||||
|
};
|
||||||
|
let json = serde_json::to_string(&full).unwrap();
|
||||||
|
let back: McpRegistration = serde_json::from_str(&json).unwrap();
|
||||||
|
|
||||||
|
assert_eq!(back.token_url.as_deref(), Some("https://as.example/token"));
|
||||||
|
assert_eq!(back.resource.as_deref(), Some("https://mcp.example/mcp"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[serial]
|
||||||
|
fn expired_token_with_old_format_registration_reports_refresh_failed() {
|
||||||
|
with_temp_cache(|| {
|
||||||
|
let dir = paths::oauth_tokens_dir();
|
||||||
|
fs::create_dir_all(&dir).unwrap();
|
||||||
|
fs::write(
|
||||||
|
paths::token_file("mcp_legacyref"),
|
||||||
|
r#"{"access_token":"stale","refresh_token":"refresh-abc","expires_at":0}"#,
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
fs::write(
|
||||||
|
dir.join("mcp_legacyref_registration.json"),
|
||||||
|
r#"{"client_id":"legacy-id"}"#,
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let rt = tokio::runtime::Builder::new_current_thread()
|
||||||
|
.enable_all()
|
||||||
|
.build()
|
||||||
|
.unwrap();
|
||||||
|
let status = rt.block_on(load_or_refresh_mcp_token("legacyref"));
|
||||||
|
|
||||||
|
assert_eq!(status, McpTokenStatus::RefreshFailed);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[serial]
|
||||||
|
fn missing_token_file_reports_not_authenticated() {
|
||||||
|
with_temp_cache(|| {
|
||||||
|
let rt = tokio::runtime::Builder::new_current_thread()
|
||||||
|
.enable_all()
|
||||||
|
.build()
|
||||||
|
.unwrap();
|
||||||
|
let status = rt.block_on(load_or_refresh_mcp_token("never-authed"));
|
||||||
|
|
||||||
|
assert_eq!(status, McpTokenStatus::NotAuthenticated);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[serial]
|
||||||
|
fn force_refresh_returns_concurrently_refreshed_token() {
|
||||||
|
with_temp_cache(|| {
|
||||||
|
fs::create_dir_all(paths::oauth_tokens_dir()).unwrap();
|
||||||
|
fs::write(
|
||||||
|
paths::token_file("mcp_force-fresh"),
|
||||||
|
r#"{"access_token":"fresh-tok","refresh_token":"r","expires_at":9999999999}"#,
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let rt = tokio::runtime::Builder::new_current_thread()
|
||||||
|
.enable_all()
|
||||||
|
.build()
|
||||||
|
.unwrap();
|
||||||
|
let token = rt.block_on(force_refresh_mcp_token("force-fresh", "rejected-tok"));
|
||||||
|
|
||||||
|
assert_eq!(token.as_deref(), Some("fresh-tok"));
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[serial]
|
||||||
|
fn force_refresh_unexpired_rejected_token_attempts_real_refresh() {
|
||||||
|
with_temp_cache(|| {
|
||||||
|
fs::create_dir_all(paths::oauth_tokens_dir()).unwrap();
|
||||||
|
fs::write(
|
||||||
|
paths::token_file("mcp_force-rejected"),
|
||||||
|
r#"{"access_token":"same-tok","refresh_token":"r","expires_at":9999999999}"#,
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let rt = tokio::runtime::Builder::new_current_thread()
|
||||||
|
.enable_all()
|
||||||
|
.build()
|
||||||
|
.unwrap();
|
||||||
|
let token = rt.block_on(force_refresh_mcp_token("force-rejected", "same-tok"));
|
||||||
|
|
||||||
|
assert_eq!(token, None);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[serial]
|
||||||
|
fn force_refresh_missing_token_file_returns_none() {
|
||||||
|
with_temp_cache(|| {
|
||||||
|
let rt = tokio::runtime::Builder::new_current_thread()
|
||||||
|
.enable_all()
|
||||||
|
.build()
|
||||||
|
.unwrap();
|
||||||
|
let token = rt.block_on(force_refresh_mcp_token(
|
||||||
|
"force-never-authed",
|
||||||
|
"rejected-tok",
|
||||||
|
));
|
||||||
|
|
||||||
|
assert_eq!(token, None);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[serial]
|
||||||
|
fn force_refresh_failed_refresh_returns_none() {
|
||||||
|
with_temp_cache(|| {
|
||||||
|
fs::create_dir_all(paths::oauth_tokens_dir()).unwrap();
|
||||||
|
fs::write(
|
||||||
|
paths::token_file("mcp_force-fail"),
|
||||||
|
r#"{"access_token":"stale","refresh_token":"r","expires_at":0}"#,
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let rt = tokio::runtime::Builder::new_current_thread()
|
||||||
|
.enable_all()
|
||||||
|
.build()
|
||||||
|
.unwrap();
|
||||||
|
let token = rt.block_on(force_refresh_mcp_token("force-fail", "stale"));
|
||||||
|
|
||||||
|
assert_eq!(token, None);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[serial]
|
||||||
|
fn force_refresh_concurrent_callers_complete_without_deadlock() {
|
||||||
|
with_temp_cache(|| {
|
||||||
|
fs::create_dir_all(paths::oauth_tokens_dir()).unwrap();
|
||||||
|
fs::write(
|
||||||
|
paths::token_file("mcp_force-concurrent"),
|
||||||
|
r#"{"access_token":"same-tok","refresh_token":"r","expires_at":9999999999}"#,
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let rt = tokio::runtime::Builder::new_current_thread()
|
||||||
|
.enable_all()
|
||||||
|
.build()
|
||||||
|
.unwrap();
|
||||||
|
let (a, b) = rt.block_on(async {
|
||||||
|
tokio::join!(
|
||||||
|
force_refresh_mcp_token("force-concurrent", "same-tok"),
|
||||||
|
force_refresh_mcp_token("force-concurrent", "same-tok"),
|
||||||
|
)
|
||||||
|
});
|
||||||
|
|
||||||
|
assert_eq!(a, None);
|
||||||
|
assert_eq!(b, None);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[serial]
|
||||||
|
fn force_refresh_respects_failure_backoff() {
|
||||||
|
with_temp_cache(|| {
|
||||||
|
fs::create_dir_all(paths::oauth_tokens_dir()).unwrap();
|
||||||
|
fs::write(
|
||||||
|
paths::token_file("mcp_force-backoff"),
|
||||||
|
r#"{"access_token":"same-tok","refresh_token":"r","expires_at":9999999999}"#,
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
note_refresh_failure("force-backoff");
|
||||||
|
|
||||||
|
let rt = tokio::runtime::Builder::new_current_thread()
|
||||||
|
.enable_all()
|
||||||
|
.build()
|
||||||
|
.unwrap();
|
||||||
|
let token = rt.block_on(force_refresh_mcp_token("force-backoff", "same-tok"));
|
||||||
|
|
||||||
|
assert_eq!(token, None);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn token_status_debug_redacts_token() {
|
||||||
|
assert_eq!(
|
||||||
|
format!("{:?}", McpTokenStatus::Token("live-secret".into())),
|
||||||
|
"Token(<redacted>)"
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
format!("{:?}", McpTokenStatus::NotAuthenticated),
|
||||||
|
"NotAuthenticated"
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
format!("{:?}", McpTokenStatus::RefreshFailed),
|
||||||
|
"RefreshFailed"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn token_status_into_token_extracts_only_token_variant() {
|
||||||
|
assert_eq!(
|
||||||
|
McpTokenStatus::Token("tok".into()).into_token(),
|
||||||
|
Some("tok".to_string())
|
||||||
|
);
|
||||||
|
assert_eq!(McpTokenStatus::NotAuthenticated.into_token(), None);
|
||||||
|
assert_eq!(McpTokenStatus::RefreshFailed.into_token(), None);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn refresh_failure_backoff_memoizes_per_server() {
|
||||||
|
assert!(!in_refresh_failure_backoff("backoff-test-server"));
|
||||||
|
|
||||||
|
note_refresh_failure("backoff-test-server");
|
||||||
|
|
||||||
|
assert!(in_refresh_failure_backoff("backoff-test-server"));
|
||||||
|
assert!(!in_refresh_failure_backoff("backoff-other-server"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn redact_refresh_error_strips_response_body() {
|
||||||
|
let with_body = anyhow!(
|
||||||
|
"Missing access_token in refresh response: {}",
|
||||||
|
r#"{"access_token":"live-secret"}"#
|
||||||
|
);
|
||||||
|
let without_body = anyhow!("no refresh token stored");
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
redact_refresh_error(&with_body),
|
||||||
|
"Missing access_token in refresh response: <response body redacted>"
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
redact_refresh_error(&without_body),
|
||||||
|
"no refresh token stored"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn cached_redirect_port_matches() {
|
fn cached_redirect_port_matches() {
|
||||||
let port = cached_redirect_port("http://127.0.0.1:49152/callback", "127.0.0.1", None);
|
let port = cached_redirect_port("http://127.0.0.1:49152/callback", "127.0.0.1", None);
|
||||||
|
|||||||
+68
-3
@@ -1,11 +1,11 @@
|
|||||||
use crate::function::{FunctionDeclaration, JsonSchema};
|
use crate::function::{self, FunctionDeclaration, JsonSchema};
|
||||||
use anyhow::{Context, Result, bail};
|
use anyhow::{Context, Result, bail};
|
||||||
use argc::{ChoiceValue, CommandValue, FlagOptionValue};
|
use argc::{ChoiceValue, CommandValue, FlagOptionValue};
|
||||||
use indexmap::IndexMap;
|
use indexmap::IndexMap;
|
||||||
|
use std::env;
|
||||||
use std::fs::File;
|
use std::fs::File;
|
||||||
use std::io::Read;
|
use std::io::Read;
|
||||||
use std::path::Path;
|
use std::path::Path;
|
||||||
use std::{env, fs};
|
|
||||||
|
|
||||||
pub fn generate_bash_declarations(
|
pub fn generate_bash_declarations(
|
||||||
mut tool_file: File,
|
mut tool_file: File,
|
||||||
@@ -23,7 +23,8 @@ pub fn generate_bash_declarations(
|
|||||||
"",
|
"",
|
||||||
env::var("TERM_WIDTH").ok().and_then(|v| v.parse().ok()),
|
env::var("TERM_WIDTH").ok().and_then(|v| v.parse().ok()),
|
||||||
)?;
|
)?;
|
||||||
fs::write(tools_file_path, &build_script)
|
let build_script = allow_empty_required_values(&build_script);
|
||||||
|
function::write_file_atomic(tools_file_path, &build_script, Some(0o755))
|
||||||
.with_context(|| format!("Failed to write built script to '{tools_file_path:?}'"))?;
|
.with_context(|| format!("Failed to write built script to '{tools_file_path:?}'"))?;
|
||||||
|
|
||||||
let command_value = argc::export(&build_script, file_name)
|
let command_value = argc::export(&build_script, file_name)
|
||||||
@@ -74,6 +75,20 @@ fn underscore(s: &str) -> String {
|
|||||||
s.replace('-', "_")
|
s.replace('-', "_")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// argc's generated required-param check uses `-z "${!name:-}"`, which
|
||||||
|
/// conflates "not provided" with "provided but empty", so a required option
|
||||||
|
/// passed an explicit empty string (e.g. `fs_write --content=''` to create an
|
||||||
|
/// empty file) is rejected as "required arguments were not provided". The
|
||||||
|
/// JSON schema we advertise to models treats `required` as *presence*, so
|
||||||
|
/// rewrite the check to a set-ness test to keep runtime behavior consistent
|
||||||
|
/// with the schema. Applied post-build so it survives every regeneration.
|
||||||
|
fn allow_empty_required_values(build_script: &str) -> String {
|
||||||
|
build_script.replace(
|
||||||
|
r#"if [[ -z "${!name:-}" ]]; then"#,
|
||||||
|
r#"if [[ -z "${!name+x}" ]]; then"#,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
fn schema_ty(t: &str) -> JsonSchema {
|
fn schema_ty(t: &str) -> JsonSchema {
|
||||||
JsonSchema {
|
JsonSchema {
|
||||||
type_value: Some(t.to_string()),
|
type_value: Some(t.to_string()),
|
||||||
@@ -147,3 +162,53 @@ fn parse_parameters_schema(flags: &[FlagOptionValue]) -> JsonSchema {
|
|||||||
required: Some(required),
|
required: Some(required),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn rewrites_argc_required_check_to_setness_test() {
|
||||||
|
let src = "# @describe test tool\n# @option --content! The content\nmain() { :; }\n";
|
||||||
|
let built = argc::build(src, "", None).expect("argc build failed");
|
||||||
|
assert!(
|
||||||
|
built.contains(r#"if [[ -z "${!name:-}" ]]; then"#),
|
||||||
|
"argc changed its generated required-param template; update allow_empty_required_values()"
|
||||||
|
);
|
||||||
|
|
||||||
|
let fixed = allow_empty_required_values(&built);
|
||||||
|
|
||||||
|
assert!(fixed.contains(r#"if [[ -z "${!name+x}" ]]; then"#));
|
||||||
|
assert!(!fixed.contains(r#"if [[ -z "${!name:-}" ]]; then"#));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(unix)]
|
||||||
|
#[test]
|
||||||
|
fn second_build_does_not_rewrite_tool_file_test() {
|
||||||
|
use std::fs;
|
||||||
|
use std::os::unix::fs::{MetadataExt, PermissionsExt};
|
||||||
|
|
||||||
|
let src = "# @describe test tool\n# @option --content! The content\nmain() { :; }\n";
|
||||||
|
let path = crate::utils::temp_file("bash-build", ".sh");
|
||||||
|
fs::write(&path, src).expect("failed to write temp script");
|
||||||
|
|
||||||
|
let declarations = generate_bash_declarations(File::open(&path).unwrap(), &path, "test")
|
||||||
|
.expect("first build failed");
|
||||||
|
assert!(!declarations.is_empty());
|
||||||
|
let first_ino = fs::metadata(&path).unwrap().ino();
|
||||||
|
|
||||||
|
let declarations = generate_bash_declarations(File::open(&path).unwrap(), &path, "test")
|
||||||
|
.expect("second build failed");
|
||||||
|
assert!(!declarations.is_empty());
|
||||||
|
|
||||||
|
let metadata = fs::metadata(&path).unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
metadata.ino(),
|
||||||
|
first_ino,
|
||||||
|
"second build rewrote an already-built tool file"
|
||||||
|
);
|
||||||
|
assert_eq!(metadata.permissions().mode() & 0o777, 0o755);
|
||||||
|
|
||||||
|
fs::remove_file(&path).unwrap();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
+37
-37
@@ -151,7 +151,7 @@ impl Rag {
|
|||||||
println!("⚙ Initializing RAG...");
|
println!("⚙ Initializing RAG...");
|
||||||
let mut data = Self::resolve_init_data(app, config)?;
|
let mut data = Self::resolve_init_data(app, config)?;
|
||||||
data.driver = config.driver.clone().unwrap_or_else(|| "yaml".to_string());
|
data.driver = config.driver.clone().unwrap_or_else(|| "yaml".to_string());
|
||||||
let mut rag = Self::create(app, name, save_path, data)?;
|
let mut rag = Self::create(app, name, save_path, data).await?;
|
||||||
let loaders = app.document_loaders.clone();
|
let loaders = app.document_loaders.clone();
|
||||||
let (spinner, spinner_rx) = Spinner::create("");
|
let (spinner, spinner_rx) = Spinner::create("");
|
||||||
abortable_run_with_spinner_rx(
|
abortable_run_with_spinner_rx(
|
||||||
@@ -298,7 +298,7 @@ impl Rag {
|
|||||||
},
|
},
|
||||||
);
|
);
|
||||||
data.driver = driver;
|
data.driver = driver;
|
||||||
let mut rag = Self::create(app, name, save_path, data)?;
|
let mut rag = Self::create(app, name, save_path, data).await?;
|
||||||
let mut paths = doc_paths.to_vec();
|
let mut paths = doc_paths.to_vec();
|
||||||
if paths.is_empty() {
|
if paths.is_empty() {
|
||||||
paths = add_documents()?;
|
paths = add_documents()?;
|
||||||
@@ -317,12 +317,12 @@ impl Rag {
|
|||||||
Ok(rag)
|
Ok(rag)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn load(app: &AppConfig, name: &str, path: &Path) -> Result<Self> {
|
pub async fn load(app: &AppConfig, name: &str, path: &Path) -> Result<Self> {
|
||||||
let err = || format!("Failed to load rag '{name}' at '{}'", path.display());
|
let err = || format!("Failed to load rag '{name}' at '{}'", path.display());
|
||||||
let content = fs::read_to_string(path).with_context(err)?;
|
let content = fs::read_to_string(path).with_context(err)?;
|
||||||
let data: RagData = serde_yaml::from_str(&content).with_context(err)?;
|
let data: RagData = serde_yaml::from_str(&content).with_context(err)?;
|
||||||
data.validate().with_context(err)?;
|
data.validate().with_context(err)?;
|
||||||
Self::create(app, name, path, data)
|
Self::create(app, name, path, data).await
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Loads a RAG from a YAML file. External drivers need an async constructor
|
/// Loads a RAG from a YAML file. External drivers need an async constructor
|
||||||
@@ -372,7 +372,7 @@ impl Rag {
|
|||||||
last_sources: RwLock::new(None),
|
last_sources: RwLock::new(None),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
_ => Self::load(app, name, path),
|
_ => Self::load(app, name, path).await,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -537,14 +537,22 @@ impl Rag {
|
|||||||
Ok(rag)
|
Ok(rag)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn create(app: &AppConfig, name: &str, path: &Path, mut data: RagData) -> Result<Self> {
|
pub async fn create(
|
||||||
|
app: &AppConfig,
|
||||||
|
name: &str,
|
||||||
|
path: &Path,
|
||||||
|
mut data: RagData,
|
||||||
|
) -> Result<Self> {
|
||||||
// Deliberately does NOT call rebuild_indexes: both callers construct the Rag
|
// Deliberately does NOT call rebuild_indexes: both callers construct the Rag
|
||||||
// before any documents are added, so rebuilding empty data would be a no-op.
|
// before any documents are added, so rebuilding empty data would be a no-op.
|
||||||
// Actual population happens later via sync_documents.
|
// Actual population happens later via sync_documents.
|
||||||
let (provider, bm25): (Box<dyn RagProvider>, _) = match data.driver.as_str() {
|
let (provider, bm25): (Box<dyn RagProvider>, _) = match data.driver.as_str() {
|
||||||
"duckdb" => {
|
"duckdb" => {
|
||||||
let db_path = providers::duckdb_path_from_yaml(path);
|
let db_path = providers::duckdb_path_from_yaml(path);
|
||||||
let dim = embedding_dim_for_model(&data.embedding_model);
|
let dim = match DuckDbProvider::introspect_dim(&db_path)? {
|
||||||
|
Some(existing) => existing,
|
||||||
|
None => probe_embedding_dim(app, &data.embedding_model).await?,
|
||||||
|
};
|
||||||
let duck = DuckDbProvider::open(&db_path, dim)?;
|
let duck = DuckDbProvider::open(&db_path, dim)?;
|
||||||
// HYDRATE — mandatory, not an optimization. The YAML file for a duckdb
|
// HYDRATE — mandatory, not an optimization. The YAML file for a duckdb
|
||||||
// RAG deliberately omits `vectors`, so `data.vectors` arrives empty from
|
// RAG deliberately omits `vectors`, so `data.vectors` arrives empty from
|
||||||
@@ -2119,21 +2127,25 @@ fn reciprocal_rank_fusion(
|
|||||||
.collect()
|
.collect()
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Map an embedding model id to its vector dimension.
|
async fn probe_embedding_dim(app: &AppConfig, model_id: &str) -> Result<usize> {
|
||||||
///
|
let model = Model::retrieve_model(app, model_id, ModelType::Embedding)?;
|
||||||
/// The DuckDB `FLOAT[N]` column type and its HNSW index are fixed at schema-creation
|
let client = init_client(&Arc::new(app.clone()), model)?;
|
||||||
/// time, so this value must be decided before the first insert. An unrecognized model
|
let out = client
|
||||||
/// falls back to 1536; if that is wrong, DuckDB raises a dimension-mismatch error on
|
.embeddings(&EmbeddingsData::new(vec!["dimension probe".into()], false))
|
||||||
/// the first insert rather than silently corrupting the schema, and the recovery is to
|
.await
|
||||||
/// delete the sidecar and re-ingest from source.
|
.with_context(|| {
|
||||||
fn embedding_dim_for_model(model_id: &str) -> usize {
|
format!(
|
||||||
match model_id {
|
"Failed to probe the embedding dimension of model '{model_id}'. \
|
||||||
m if m.contains("3-large") => 3072,
|
Creating a duckdb RAG requires one call to the embedding endpoint."
|
||||||
m if m.contains("3-small") || m.contains("ada-002") => 1536,
|
)
|
||||||
m if m.contains("nomic-embed-text") || m.contains("all-minilm") => 768,
|
})?;
|
||||||
m if m.contains("jina-embeddings-v2") => 1024,
|
let dim = out.first().map(|v| v.len()).unwrap_or(0);
|
||||||
_ => 1536,
|
|
||||||
|
if dim == 0 {
|
||||||
|
bail!("Embedding model '{model_id}' returned an empty vector during the dimension probe");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
Ok(dim)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// True only for "the vault does not hold this key".
|
/// True only for "the vault does not hold this key".
|
||||||
@@ -2673,18 +2685,6 @@ mod tests {
|
|||||||
assert_eq!(data.attached_source_label(), "[external collection]");
|
assert_eq!(data.attached_source_label(), "[external collection]");
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn embedding_dim_for_model_maps_known_models() {
|
|
||||||
assert_eq!(embedding_dim_for_model("text-embedding-3-large"), 3072);
|
|
||||||
assert_eq!(embedding_dim_for_model("text-embedding-3-small"), 1536);
|
|
||||||
assert_eq!(embedding_dim_for_model("text-embedding-ada-002"), 1536);
|
|
||||||
assert_eq!(embedding_dim_for_model("nomic-embed-text"), 768);
|
|
||||||
assert_eq!(embedding_dim_for_model("all-minilm"), 768);
|
|
||||||
assert_eq!(embedding_dim_for_model("jina-embeddings-v2-base-en"), 1024);
|
|
||||||
// Unknown models fall back to the OpenAI-compatible default.
|
|
||||||
assert_eq!(embedding_dim_for_model("some-unknown-model"), 1536);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn document_id_round_trip() {
|
fn document_id_round_trip() {
|
||||||
let id = DocumentId::new(5, 17);
|
let id = DocumentId::new(5, 17);
|
||||||
@@ -3044,7 +3044,7 @@ mod tests {
|
|||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn reciprocal_rank_fusion_empty_lists() {
|
fn reciprocal_rank_fusion_empty_lists() {
|
||||||
let result = super::reciprocal_rank_fusion(vec![], vec![], 5);
|
let result = reciprocal_rank_fusion(vec![], vec![], 5);
|
||||||
assert!(result.is_empty(), "empty input should produce empty output");
|
assert!(result.is_empty(), "empty input should produce empty output");
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -3052,7 +3052,7 @@ mod tests {
|
|||||||
fn reciprocal_rank_fusion_deduplicates_across_signals() {
|
fn reciprocal_rank_fusion_deduplicates_across_signals() {
|
||||||
let doc_a = DocumentId::new(0, 0);
|
let doc_a = DocumentId::new(0, 0);
|
||||||
let doc_b = DocumentId::new(0, 1);
|
let doc_b = DocumentId::new(0, 1);
|
||||||
let result = super::reciprocal_rank_fusion(
|
let result = reciprocal_rank_fusion(
|
||||||
vec![vec![doc_a, doc_b], vec![doc_a, doc_b]],
|
vec![vec![doc_a, doc_b], vec![doc_a, doc_b]],
|
||||||
vec![1.0, 1.0],
|
vec![1.0, 1.0],
|
||||||
5,
|
5,
|
||||||
@@ -3069,7 +3069,7 @@ mod tests {
|
|||||||
#[test]
|
#[test]
|
||||||
fn reciprocal_rank_fusion_respects_top_k() {
|
fn reciprocal_rank_fusion_respects_top_k() {
|
||||||
let docs: Vec<DocumentId> = (0..10).map(|i| DocumentId::new(0, i)).collect();
|
let docs: Vec<DocumentId> = (0..10).map(|i| DocumentId::new(0, i)).collect();
|
||||||
let result = super::reciprocal_rank_fusion(vec![docs], vec![1.0], 3);
|
let result = reciprocal_rank_fusion(vec![docs], vec![1.0], 3);
|
||||||
assert_eq!(result.len(), 3, "result should be capped at top_k=3");
|
assert_eq!(result.len(), 3, "result should be capped at top_k=3");
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -3077,7 +3077,7 @@ mod tests {
|
|||||||
fn reciprocal_rank_fusion_weights_affect_ranking() {
|
fn reciprocal_rank_fusion_weights_affect_ranking() {
|
||||||
let doc_a = DocumentId::new(0, 0);
|
let doc_a = DocumentId::new(0, 0);
|
||||||
let doc_b = DocumentId::new(0, 1);
|
let doc_b = DocumentId::new(0, 1);
|
||||||
let result = super::reciprocal_rank_fusion(
|
let result = reciprocal_rank_fusion(
|
||||||
vec![vec![doc_a, doc_b], vec![doc_b, doc_a]],
|
vec![vec![doc_a, doc_b], vec![doc_b, doc_a]],
|
||||||
vec![10.0, 1.0],
|
vec![10.0, 1.0],
|
||||||
2,
|
2,
|
||||||
|
|||||||
+211
-9
@@ -5,7 +5,7 @@ use std::collections::HashMap;
|
|||||||
use anyhow::{Context, Result, anyhow, bail};
|
use anyhow::{Context, Result, anyhow, bail};
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use duckdb::types::Value;
|
use duckdb::types::Value;
|
||||||
use duckdb::{AccessMode, Config, Connection};
|
use duckdb::{AccessMode, Config, Connection, OptionalExt};
|
||||||
use indexmap::IndexMap;
|
use indexmap::IndexMap;
|
||||||
use log::warn;
|
use log::warn;
|
||||||
use std::path::{Path, PathBuf};
|
use std::path::{Path, PathBuf};
|
||||||
@@ -30,6 +30,11 @@ pub(crate) fn duckdb_path_from_yaml(yaml_path: &Path) -> PathBuf {
|
|||||||
yaml_path.with_extension("duckdb")
|
yaml_path.with_extension("duckdb")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn parse_float_array_dim(data_type: &str) -> Option<usize> {
|
||||||
|
let inner = data_type.strip_prefix("FLOAT[")?.strip_suffix(']')?;
|
||||||
|
inner.parse::<usize>().ok().filter(|&n| n > 0)
|
||||||
|
}
|
||||||
|
|
||||||
/// The shared connection together with the access mode it was opened with.
|
/// The shared connection together with the access mode it was opened with.
|
||||||
///
|
///
|
||||||
/// `conn` is an `Option` only so that an upgrade can DROP the read-only connection
|
/// `conn` is an `Option` only so that an upgrade can DROP the read-only connection
|
||||||
@@ -42,6 +47,11 @@ struct ConnHandle {
|
|||||||
/// `duplicate()` clone sharing the `Arc` observes an upgrade performed through any
|
/// `duplicate()` clone sharing the `Arc` observes an upgrade performed through any
|
||||||
/// other handle instead of keeping its own stale copy of the mode.
|
/// other handle instead of keeping its own stale copy of the mode.
|
||||||
writable: bool,
|
writable: bool,
|
||||||
|
/// Embedding dimension the `FLOAT[N]` column was opened (or last rebuilt) with.
|
||||||
|
/// Shared for the same reason as `writable`: a self-healing rebuild through one
|
||||||
|
/// handle updates the width, and every `duplicate()` clone must cast with the new
|
||||||
|
/// width instead of erroring on a healthy store with its stale copy.
|
||||||
|
dim: usize,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl ConnHandle {
|
impl ConnHandle {
|
||||||
@@ -68,8 +78,6 @@ impl ConnHandle {
|
|||||||
pub struct DuckDbProvider {
|
pub struct DuckDbProvider {
|
||||||
path: PathBuf,
|
path: PathBuf,
|
||||||
conn: Arc<Mutex<ConnHandle>>,
|
conn: Arc<Mutex<ConnHandle>>,
|
||||||
/// Embedding dimension; fixed at open time because the `FLOAT[N]` column depends on it.
|
|
||||||
dim: usize,
|
|
||||||
/// True once an FTS index has been built on `documents`. Until then
|
/// True once an FTS index has been built on `documents`. Until then
|
||||||
/// `fts_main_documents.match_bm25` does not exist and any keyword query would
|
/// `fts_main_documents.match_bm25` does not exist and any keyword query would
|
||||||
/// fail with a DuckDB catalog error. Backs `has_native_keyword_search`.
|
/// fail with a DuckDB catalog error. Backs `has_native_keyword_search`.
|
||||||
@@ -96,8 +104,8 @@ impl DuckDbProvider {
|
|||||||
conn: Arc::new(Mutex::new(ConnHandle {
|
conn: Arc::new(Mutex::new(ConnHandle {
|
||||||
conn: Some(conn),
|
conn: Some(conn),
|
||||||
writable,
|
writable,
|
||||||
|
dim,
|
||||||
})),
|
})),
|
||||||
dim,
|
|
||||||
fts_ready: AtomicBool::new(fts_exists),
|
fts_ready: AtomicBool::new(fts_exists),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
@@ -164,6 +172,35 @@ impl DuckDbProvider {
|
|||||||
Ok(conn)
|
Ok(conn)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn introspect_dim(db_path: &Path) -> Result<Option<usize>> {
|
||||||
|
if !db_path.exists() {
|
||||||
|
return Ok(None);
|
||||||
|
}
|
||||||
|
let conn = Self::open_read_only(db_path).with_context(|| {
|
||||||
|
format!(
|
||||||
|
"Cannot inspect the existing RAG store at '{}'",
|
||||||
|
db_path.display()
|
||||||
|
)
|
||||||
|
})?;
|
||||||
|
let ty: Option<String> = conn
|
||||||
|
.query_row(
|
||||||
|
"SELECT data_type FROM duckdb_columns() \
|
||||||
|
WHERE table_name = 'vectors' AND column_name = 'embedding'",
|
||||||
|
[],
|
||||||
|
|r| r.get(0),
|
||||||
|
)
|
||||||
|
.optional()
|
||||||
|
.with_context(|| {
|
||||||
|
format!(
|
||||||
|
"Failed to introspect the embedding dimension of the DuckDB store \
|
||||||
|
at '{}'",
|
||||||
|
db_path.display()
|
||||||
|
)
|
||||||
|
})?;
|
||||||
|
|
||||||
|
Ok(ty.and_then(|t| parse_float_array_dim(&t)))
|
||||||
|
}
|
||||||
|
|
||||||
/// Open the store read-write and make sure its schema exists. Exactly one process
|
/// Open the store read-write and make sure its schema exists. Exactly one process
|
||||||
/// may hold such a handle, and no reader from another process may hold it meanwhile.
|
/// may hold such a handle, and no reader from another process may hold it meanwhile.
|
||||||
fn open_read_write(db_path: &Path, dim: usize) -> Result<Connection> {
|
fn open_read_write(db_path: &Path, dim: usize) -> Result<Connection> {
|
||||||
@@ -245,7 +282,7 @@ impl DuckDbProvider {
|
|||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
drop(handle.conn.take());
|
drop(handle.conn.take());
|
||||||
match Self::open_read_write(&self.path, self.dim) {
|
match Self::open_read_write(&self.path, handle.dim) {
|
||||||
Ok(conn) => {
|
Ok(conn) => {
|
||||||
handle.conn = Some(conn);
|
handle.conn = Some(conn);
|
||||||
handle.writable = true;
|
handle.writable = true;
|
||||||
@@ -428,13 +465,33 @@ impl RagProvider for DuckDbProvider {
|
|||||||
if embedding.iter().any(|f| !f.is_finite()) {
|
if embedding.iter().any(|f| !f.is_finite()) {
|
||||||
bail!("Query embedding contains a non-finite value (NaN or infinity)");
|
bail!("Query embedding contains a non-finite value (NaN or infinity)");
|
||||||
}
|
}
|
||||||
|
let handle = self.lock_conn()?;
|
||||||
|
let dim = handle.dim;
|
||||||
|
if embedding.len() != dim {
|
||||||
|
let rows: i64 = handle
|
||||||
|
.conn()?
|
||||||
|
.query_row("SELECT count(*) FROM vectors", [], |r| r.get(0))
|
||||||
|
.context("Failed to count vectors before a dimension-mismatch query")?;
|
||||||
|
if rows == 0 {
|
||||||
|
// A never-synced store answers "nothing", not a cast error.
|
||||||
|
return Ok(Vec::new());
|
||||||
|
}
|
||||||
|
|
||||||
|
bail!(
|
||||||
|
"RAG store at '{}' was built with {dim}-dim embeddings, but the \
|
||||||
|
embedding model now returns {}-dim vectors. The embedding model \
|
||||||
|
changed since ingestion. Re-embed the documents, or delete the \
|
||||||
|
sidecar file and re-ingest.",
|
||||||
|
self.path.display(),
|
||||||
|
embedding.len()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
let vals: String = embedding
|
let vals: String = embedding
|
||||||
.iter()
|
.iter()
|
||||||
.map(|f| f.to_string())
|
.map(|f| f.to_string())
|
||||||
.collect::<Vec<_>>()
|
.collect::<Vec<_>>()
|
||||||
.join(", ");
|
.join(", ");
|
||||||
let dim = self.dim;
|
|
||||||
let handle = self.lock_conn()?;
|
|
||||||
// array_cosine_distance requires a FLOAT[N] ARRAY, not the LIST type FLOAT[].
|
// array_cosine_distance requires a FLOAT[N] ARRAY, not the LIST type FLOAT[].
|
||||||
// ORDER BY distance ASC is required for the planner to use hnsw_idx; the
|
// ORDER BY distance ASC is required for the planner to use hnsw_idx; the
|
||||||
// similarity form (DESC) does NOT trigger the ANN index. Distance is converted
|
// similarity form (DESC) does NOT trigger the ANN index. Distance is converted
|
||||||
@@ -554,7 +611,28 @@ impl RagProvider for DuckDbProvider {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
let dim = self.dim;
|
let dim = match data.vectors.first() {
|
||||||
|
None => self.lock_conn()?.dim,
|
||||||
|
Some((_, first)) => {
|
||||||
|
let dim = first.len();
|
||||||
|
if let Some((doc_id, other)) = data.vectors.iter().find(|(_, e)| e.len() != dim) {
|
||||||
|
let matching = data.vectors.values().filter(|e| e.len() == dim).count();
|
||||||
|
|
||||||
|
bail!(
|
||||||
|
"Refusing to rebuild the RAG store at '{}': the rebuild batch \
|
||||||
|
mixes {dim}-dim and {}-dim vectors ({matching} vs {} vectors; \
|
||||||
|
first mismatch: document {}). Re-embed the documents, or \
|
||||||
|
delete the sidecar file and re-ingest.",
|
||||||
|
self.path.display(),
|
||||||
|
other.len(),
|
||||||
|
data.vectors.len() - matching,
|
||||||
|
doc_id.0
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
dim
|
||||||
|
}
|
||||||
|
};
|
||||||
// THE write path. Everything above this line only reads, so the upgrade happens
|
// THE write path. Everything above this line only reads, so the upgrade happens
|
||||||
// here, after both guards have had their say: a rebuild that is going to be
|
// here, after both guards have had their say: a rebuild that is going to be
|
||||||
// refused must not first take the exclusive lock away from other processes.
|
// refused must not first take the exclusive lock away from other processes.
|
||||||
@@ -664,6 +742,8 @@ impl RagProvider for DuckDbProvider {
|
|||||||
// fall back to local BM25, which is also empty, and therefore correct.
|
// fall back to local BM25, which is also empty, and therefore correct.
|
||||||
self.fts_ready.store(doc_count > 0, Ordering::Relaxed);
|
self.fts_ready.store(doc_count > 0, Ordering::Relaxed);
|
||||||
|
|
||||||
|
handle.dim = dim;
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -729,10 +809,11 @@ impl RagProvider for DuckDbProvider {
|
|||||||
// Sharing the Arc also shares the ACCESS MODE, which lives inside the ConnHandle
|
// Sharing the Arc also shares the ACCESS MODE, which lives inside the ConnHandle
|
||||||
// rather than beside it: when one handle upgrades itself to read-write, every
|
// rather than beside it: when one handle upgrades itself to read-write, every
|
||||||
// clone is upgraded with it and none is left holding a stale "read-only" belief.
|
// clone is upgraded with it and none is left holding a stale "read-only" belief.
|
||||||
|
// The embedding dimension lives there too, so a self-healing rebuild through
|
||||||
|
// one handle updates the width every clone casts with.
|
||||||
Box::new(DuckDbProvider {
|
Box::new(DuckDbProvider {
|
||||||
path: self.path.clone(),
|
path: self.path.clone(),
|
||||||
conn: Arc::clone(&self.conn),
|
conn: Arc::clone(&self.conn),
|
||||||
dim: self.dim,
|
|
||||||
fts_ready: AtomicBool::new(self.fts_ready.load(Ordering::Relaxed)),
|
fts_ready: AtomicBool::new(self.fts_ready.load(Ordering::Relaxed)),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
@@ -1074,6 +1155,127 @@ mod tests {
|
|||||||
.expect("a fresh RAG with nothing indexed must rebuild cleanly");
|
.expect("a fresh RAG with nothing indexed must rebuild cleanly");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_float_array_dim_handles_arrays_lists_and_scalars() {
|
||||||
|
assert_eq!(parse_float_array_dim("FLOAT[768]"), Some(768));
|
||||||
|
assert_eq!(parse_float_array_dim("FLOAT[]"), None);
|
||||||
|
assert_eq!(parse_float_array_dim("VARCHAR"), None);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn introspect_dim_round_trips_the_open_dim() {
|
||||||
|
let db = TempDb::new("introspect");
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
DuckDbProvider::introspect_dim(&db.path).unwrap(),
|
||||||
|
None,
|
||||||
|
"a file that does not exist has no dim"
|
||||||
|
);
|
||||||
|
{
|
||||||
|
let _provider = DuckDbProvider::open(&db.path, 5).unwrap();
|
||||||
|
}
|
||||||
|
assert_eq!(DuckDbProvider::introspect_dim(&db.path).unwrap(), Some(5));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn introspect_dim_propagates_an_unopenable_existing_file() {
|
||||||
|
let db = TempDb::new("introspectgarbage");
|
||||||
|
fs::write(&db.path, b"not a duckdb database").unwrap();
|
||||||
|
|
||||||
|
let err = DuckDbProvider::introspect_dim(&db.path).unwrap_err();
|
||||||
|
|
||||||
|
assert!(
|
||||||
|
format!("{err:#}").contains(&format!(
|
||||||
|
"Cannot inspect the existing RAG store at '{}'",
|
||||||
|
db.path.display()
|
||||||
|
)),
|
||||||
|
"an existing-but-unopenable file must be an error naming the path; got: {err:#}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn rebuild_indexes_self_heals_dim_from_the_vectors_it_writes() {
|
||||||
|
let db = TempDb::new("selfheal");
|
||||||
|
let mut provider = DuckDbProvider::open(&db.path, 5).unwrap();
|
||||||
|
let mut data = minimal_rag_data();
|
||||||
|
data.vectors.insert(DocumentId(0), vec![0.1, 0.2, 0.3]);
|
||||||
|
|
||||||
|
provider.rebuild_indexes(&data, true).await.unwrap();
|
||||||
|
|
||||||
|
let results = provider
|
||||||
|
.vector_search(&[0.1, 0.2, 0.3], 5, 0.0)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(results.len(), 1);
|
||||||
|
|
||||||
|
drop(provider);
|
||||||
|
assert_eq!(DuckDbProvider::introspect_dim(&db.path).unwrap(), Some(3));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn a_self_healed_dim_is_visible_through_duplicate_clones() {
|
||||||
|
let db = TempDb::new("dimdup");
|
||||||
|
let mut provider = DuckDbProvider::open(&db.path, 5).unwrap();
|
||||||
|
let dup = provider.duplicate(&minimal_rag_data());
|
||||||
|
|
||||||
|
let mut data = minimal_rag_data();
|
||||||
|
data.vectors.insert(DocumentId(0), vec![0.1, 0.2, 0.3]);
|
||||||
|
provider.rebuild_indexes(&data, true).await.unwrap();
|
||||||
|
|
||||||
|
let results = dup.vector_search(&[0.1, 0.2, 0.3], 5, 0.0).await.unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
results.len(),
|
||||||
|
1,
|
||||||
|
"a duplicate() clone must observe the dim written by a rebuild through \
|
||||||
|
the original"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn rebuild_indexes_rejects_mixed_dim_vectors() {
|
||||||
|
let db = TempDb::new("mixeddim");
|
||||||
|
let mut provider = DuckDbProvider::open(&db.path, 3).unwrap();
|
||||||
|
let mut data = minimal_rag_data();
|
||||||
|
data.vectors.insert(DocumentId(0), vec![0.1, 0.2, 0.3]);
|
||||||
|
data.vectors.insert(DocumentId(1), vec![0.1, 0.2, 0.3, 0.4]);
|
||||||
|
|
||||||
|
let err = provider.rebuild_indexes(&data, true).await.unwrap_err();
|
||||||
|
|
||||||
|
assert!(
|
||||||
|
err.to_string().contains("rebuild batch mixes"),
|
||||||
|
"got: {err}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn vector_search_dim_mismatch_on_empty_store_returns_nothing() {
|
||||||
|
let db = TempDb::new("dimempty");
|
||||||
|
let provider = DuckDbProvider::open(&db.path, 3).unwrap();
|
||||||
|
|
||||||
|
let results = provider.vector_search(&[0.1, 0.2], 5, 0.0).await.unwrap();
|
||||||
|
|
||||||
|
assert!(
|
||||||
|
results.is_empty(),
|
||||||
|
"a never-synced store must answer 'nothing', not a cast error"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn vector_search_dim_mismatch_on_populated_store_errors() {
|
||||||
|
let db = TempDb::new("dimfull");
|
||||||
|
let mut provider = DuckDbProvider::open(&db.path, 3).unwrap();
|
||||||
|
let mut data = minimal_rag_data();
|
||||||
|
data.vectors.insert(DocumentId(0), vec![0.1, 0.2, 0.3]);
|
||||||
|
provider.rebuild_indexes(&data, true).await.unwrap();
|
||||||
|
|
||||||
|
let err = provider
|
||||||
|
.vector_search(&[0.1, 0.2], 5, 0.0)
|
||||||
|
.await
|
||||||
|
.unwrap_err();
|
||||||
|
|
||||||
|
assert!(err.to_string().contains("was built with"), "got: {err}");
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn duplicate_shares_the_same_connection() {
|
async fn duplicate_shares_the_same_connection() {
|
||||||
let db = TempDb::new("dup");
|
let db = TempDb::new("dup");
|
||||||
|
|||||||
+34
-2
@@ -39,8 +39,10 @@ static INLINE_CODE_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"`([^`\n]+
|
|||||||
static IMAGE_RE: LazyLock<Regex> =
|
static IMAGE_RE: LazyLock<Regex> =
|
||||||
LazyLock::new(|| Regex::new(r"!\[([^\]]*)\]\(([^)]+)\)").unwrap());
|
LazyLock::new(|| Regex::new(r"!\[([^\]]*)\]\(([^)]+)\)").unwrap());
|
||||||
static LINK_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\[([^\]]+)\]\(([^)]+)\)").unwrap());
|
static LINK_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\[([^\]]+)\]\(([^)]+)\)").unwrap());
|
||||||
static BOLD_AST_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\*\*([^*\n]+)\*\*").unwrap());
|
static BOLD_AST_RE: LazyLock<Regex> =
|
||||||
static BOLD_US_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"__([^_\n]+)__").unwrap());
|
LazyLock::new(|| Regex::new(r"\*\*((?:[^*\n]|\*(?!\*))+?)\*\*").unwrap());
|
||||||
|
static BOLD_US_RE: LazyLock<Regex> =
|
||||||
|
LazyLock::new(|| Regex::new(r"__((?:[^_\n]|_(?!_))+?)__").unwrap());
|
||||||
static ITALIC_AST_RE: LazyLock<Regex> =
|
static ITALIC_AST_RE: LazyLock<Regex> =
|
||||||
LazyLock::new(|| Regex::new(r"(?<![*\w])\*(?!\s)([^*\n]+?)(?<!\s)\*(?!\*)").unwrap());
|
LazyLock::new(|| Regex::new(r"(?<![*\w])\*(?!\s)([^*\n]+?)(?<!\s)\*(?!\*)").unwrap());
|
||||||
static ITALIC_US_RE: LazyLock<Regex> =
|
static ITALIC_US_RE: LazyLock<Regex> =
|
||||||
@@ -2419,6 +2421,36 @@ std::error::Error>> {
|
|||||||
assert!(result.contains("\x1b[9m"), "strikethrough SGR: {result:?}");
|
assert!(result.contains("\x1b[9m"), "strikethrough SGR: {result:?}");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn bold_asterisk_wraps_italic_asterisk() {
|
||||||
|
let styles = test_styles();
|
||||||
|
|
||||||
|
let result = apply_inline("**loud *soft* loud**", &styles);
|
||||||
|
|
||||||
|
assert!(
|
||||||
|
!result.contains("**"),
|
||||||
|
"outer bold markers stripped: {result:?}"
|
||||||
|
);
|
||||||
|
assert!(result.contains("\x1b[1m"), "bold SGR present: {result:?}");
|
||||||
|
assert!(result.contains("\x1b[3m"), "italic SGR present: {result:?}");
|
||||||
|
assert!(result.contains("soft"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn bold_underscore_wraps_italic_underscore() {
|
||||||
|
let styles = test_styles();
|
||||||
|
|
||||||
|
let result = apply_inline("__loud _soft_ loud__", &styles);
|
||||||
|
|
||||||
|
assert!(
|
||||||
|
!result.contains("__"),
|
||||||
|
"outer bold markers stripped: {result:?}"
|
||||||
|
);
|
||||||
|
assert!(result.contains("\x1b[1m"), "bold SGR present: {result:?}");
|
||||||
|
assert!(result.contains("\x1b[3m"), "italic SGR present: {result:?}");
|
||||||
|
assert!(result.contains("soft"));
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn bold_wraps_inline_code() {
|
fn bold_wraps_inline_code() {
|
||||||
let styles = test_styles();
|
let styles = test_styles();
|
||||||
|
|||||||
+6
-2
@@ -19,8 +19,8 @@ use crate::config::{AssetCategory, paths};
|
|||||||
use crate::function::supervisor::{GuardrailAction, check_pending_agents_guardrail};
|
use crate::function::supervisor::{GuardrailAction, check_pending_agents_guardrail};
|
||||||
use crate::render::render_error;
|
use crate::render::render_error;
|
||||||
use crate::utils::{
|
use crate::utils::{
|
||||||
AbortSignal, SHELL, abortable_run_with_spinner, create_abort_signal, dimmed_text, run_command,
|
AbortSignal, SHELL, abortable_run_with_spinner, create_abort_signal, dimmed_text,
|
||||||
set_text, temp_file,
|
drain_stale_tty_input, run_command, set_text, temp_file,
|
||||||
};
|
};
|
||||||
|
|
||||||
use crate::sandbox::SANDBOX_ENV_FLAG;
|
use crate::sandbox::SANDBOX_ENV_FLAG;
|
||||||
@@ -411,6 +411,10 @@ Type ".help" for additional help.
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Discard any stray terminal-query reply bytes (e.g. late colorsaurus
|
||||||
|
// OSC 11 / DA1 responses) so they don't get injected into the prompt.
|
||||||
|
drain_stale_tty_input();
|
||||||
|
|
||||||
loop {
|
loop {
|
||||||
if self.abort_signal.aborted_ctrld() {
|
if self.abort_signal.aborted_ctrld() {
|
||||||
break;
|
break;
|
||||||
|
|||||||
+296
-15
@@ -80,6 +80,8 @@ pub(crate) struct CredentialSpec {
|
|||||||
pub env_var: String,
|
pub env_var: String,
|
||||||
pub proxy_managed: bool,
|
pub proxy_managed: bool,
|
||||||
pub inject: Vec<InjectRule>,
|
pub inject: Vec<InjectRule>,
|
||||||
|
pub custom_hosts: Vec<String>,
|
||||||
|
pub custom_allow_entries: Vec<String>,
|
||||||
pub servers: Vec<String>,
|
pub servers: Vec<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -110,8 +112,13 @@ pub(crate) fn collect_credentials(
|
|||||||
}
|
}
|
||||||
|
|
||||||
let mut by_secret: BTreeMap<String, Aggregate> = BTreeMap::new();
|
let mut by_secret: BTreeMap<String, Aggregate> = BTreeMap::new();
|
||||||
|
let mut hosts_by_server: BTreeMap<String, ServerSecretHosts> = BTreeMap::new();
|
||||||
for (server_name, config) in servers {
|
for (server_name, config) in servers {
|
||||||
for occurrence in collect_server_occurrences(config)? {
|
let occurrences = collect_server_occurrences(config)?;
|
||||||
|
if !occurrences.is_empty() {
|
||||||
|
hosts_by_server.insert(server_name.clone(), server_secret_hosts(config));
|
||||||
|
}
|
||||||
|
for occurrence in occurrences {
|
||||||
let agg = by_secret
|
let agg = by_secret
|
||||||
.entry(occurrence.secret_name)
|
.entry(occurrence.secret_name)
|
||||||
.or_insert_with(|| Aggregate {
|
.or_insert_with(|| Aggregate {
|
||||||
@@ -152,8 +159,9 @@ pub(crate) fn collect_credentials(
|
|||||||
eprintln!(
|
eprintln!(
|
||||||
"MCP secrets {} all target the '{header}' header for '{domain}'. The \
|
"MCP secrets {} all target the '{header}' header for '{domain}'. The \
|
||||||
sandbox proxy cannot tell which one a given request needs, so these \
|
sandbox proxy cannot tell which one a given request needs, so these \
|
||||||
secrets will be resolved from environment variables inside the \
|
secrets will be provisioned as placeholder-based custom secrets \
|
||||||
sandbox instead.",
|
instead: each env var holds a unique placeholder that the proxy \
|
||||||
|
swaps for the real value in outbound request headers.",
|
||||||
quoted_list(secrets)
|
quoted_list(secrets)
|
||||||
);
|
);
|
||||||
slot_conflicted.extend(secrets.iter().cloned());
|
slot_conflicted.extend(secrets.iter().cloned());
|
||||||
@@ -191,6 +199,21 @@ pub(crate) fn collect_credentials(
|
|||||||
}
|
}
|
||||||
|
|
||||||
let proxy_managed = agg.all_injectable && !slot_conflicted.contains(&secret_name);
|
let proxy_managed = agg.all_injectable && !slot_conflicted.contains(&secret_name);
|
||||||
|
let (custom_hosts, custom_allow_entries) = if proxy_managed {
|
||||||
|
(Vec::new(), Vec::new())
|
||||||
|
} else {
|
||||||
|
let mut targets: BTreeSet<String> = BTreeSet::new();
|
||||||
|
let mut allow: BTreeSet<String> = BTreeSet::new();
|
||||||
|
for hosts in agg
|
||||||
|
.servers
|
||||||
|
.iter()
|
||||||
|
.filter_map(|server| hosts_by_server.get(server))
|
||||||
|
{
|
||||||
|
targets.extend(hosts.targets.iter().cloned());
|
||||||
|
allow.extend(hosts.allow_entries.iter().cloned());
|
||||||
|
}
|
||||||
|
(targets.into_iter().collect(), allow.into_iter().collect())
|
||||||
|
};
|
||||||
credentials.push(CredentialSpec {
|
credentials.push(CredentialSpec {
|
||||||
secret_name,
|
secret_name,
|
||||||
service_id,
|
service_id,
|
||||||
@@ -201,6 +224,8 @@ pub(crate) fn collect_credentials(
|
|||||||
} else {
|
} else {
|
||||||
Vec::new()
|
Vec::new()
|
||||||
},
|
},
|
||||||
|
custom_hosts,
|
||||||
|
custom_allow_entries,
|
||||||
servers: agg.servers.into_iter().collect(),
|
servers: agg.servers.into_iter().collect(),
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
@@ -365,6 +390,98 @@ fn parse_https_domain(raw: &str) -> Option<String> {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Default, PartialEq, Eq)]
|
||||||
|
pub(crate) struct ServerSecretHosts {
|
||||||
|
pub targets: BTreeSet<String>,
|
||||||
|
pub allow_entries: BTreeSet<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn server_secret_hosts(server: &Value) -> ServerSecretHosts {
|
||||||
|
let mut hosts = ServerSecretHosts::default();
|
||||||
|
scrape_hosts_value(server, &mut hosts);
|
||||||
|
|
||||||
|
if let Some(host) = server
|
||||||
|
.get("url")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.and_then(|raw| Url::parse(raw).ok())
|
||||||
|
.and_then(|url| url.host_str().map(str::to_string))
|
||||||
|
.filter(|host| is_usable_host(host))
|
||||||
|
{
|
||||||
|
hosts.targets.insert(host);
|
||||||
|
}
|
||||||
|
|
||||||
|
hosts
|
||||||
|
}
|
||||||
|
|
||||||
|
fn scrape_hosts_value(value: &Value, out: &mut ServerSecretHosts) {
|
||||||
|
match value {
|
||||||
|
Value::String(s) => {
|
||||||
|
for url in urls_in_text(s) {
|
||||||
|
let Some(host) = url.host_str().filter(|h| is_usable_host(h)) else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
out.targets.insert(host.to_string());
|
||||||
|
if let Some(entry) = allow_entry_for_parsed(&url) {
|
||||||
|
out.allow_entries.insert(entry);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Value::Object(map) => {
|
||||||
|
for v in map.values() {
|
||||||
|
scrape_hosts_value(v, out);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Value::Array(arr) => {
|
||||||
|
for v in arr {
|
||||||
|
scrape_hosts_value(v, out);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A host usable in the sbx target and allow grammars: non-empty, not a
|
||||||
|
/// bracketed IPv6 literal (no spelling in either grammar), and not an
|
||||||
|
/// unresolved `{{placeholder}}` fragment (would register garbage targets
|
||||||
|
/// and could fail kit validation at create).
|
||||||
|
fn is_usable_host(host: &str) -> bool {
|
||||||
|
!host.is_empty() && !host.starts_with('[') && !host.contains('{') && !host.contains('}')
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Every http(s) URL embedded in `text`. Tokens end at whitespace, quotes,
|
||||||
|
/// or URL-hostile punctuation so comma- or bracket-separated lists don't
|
||||||
|
/// bleed into one another.
|
||||||
|
fn urls_in_text(text: &str) -> Vec<Url> {
|
||||||
|
const TERMINATORS: &[char] = &['"', '\'', ',', ';', '(', ')', '[', ']', '{', '}', '<', '>'];
|
||||||
|
|
||||||
|
let mut out = Vec::new();
|
||||||
|
for (idx, _) in text.match_indices("http") {
|
||||||
|
let candidate = &text[idx..];
|
||||||
|
if !candidate.starts_with("http://") && !candidate.starts_with("https://") {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let end = candidate
|
||||||
|
.find(|c: char| c.is_whitespace() || TERMINATORS.contains(&c))
|
||||||
|
.unwrap_or(candidate.len());
|
||||||
|
if let Ok(url) = Url::parse(&candidate[..end]) {
|
||||||
|
out.push(url);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
out
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Extracts the host of every http(s) URL embedded in `text`, ports stripped.
|
||||||
|
/// Bracketed IPv6 hosts and placeholder fragments are skipped.
|
||||||
|
pub(crate) fn hosts_in_text(text: &str) -> BTreeSet<String> {
|
||||||
|
urls_in_text(text)
|
||||||
|
.iter()
|
||||||
|
.filter_map(|url| url.host_str())
|
||||||
|
.filter(|host| is_usable_host(host))
|
||||||
|
.map(str::to_string)
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
/// Collects a network allow-list entry for every remote MCP server `url`
|
/// Collects a network allow-list entry for every remote MCP server `url`
|
||||||
/// (http/https), so user-configured servers are reachable regardless of how,
|
/// (http/https), so user-configured servers are reachable regardless of how,
|
||||||
/// or whether, their credentials are provisioned. Https on the default port
|
/// or whether, their credentials are provisioned. Https on the default port
|
||||||
@@ -394,13 +511,14 @@ pub(crate) fn allow_entry_for_url(raw: &str) -> Option<String> {
|
|||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
|
|
||||||
let host = url.host_str()?;
|
allow_entry_for_parsed(&url)
|
||||||
if host.is_empty() || host.starts_with('[') {
|
}
|
||||||
return None;
|
|
||||||
}
|
fn allow_entry_for_parsed(url: &Url) -> Option<String> {
|
||||||
|
let host = url.host_str().filter(|h| is_usable_host(h))?;
|
||||||
|
|
||||||
match url.port_or_known_default() {
|
match url.port_or_known_default() {
|
||||||
Some(443) if scheme == "https" => Some(host.to_string()),
|
Some(443) if url.scheme() == "https" => Some(host.to_string()),
|
||||||
Some(port) => Some(format!("{host}:{port}")),
|
Some(port) => Some(format!("{host}:{port}")),
|
||||||
None => None,
|
None => None,
|
||||||
}
|
}
|
||||||
@@ -489,12 +607,17 @@ pub(crate) fn render_mixin_document(
|
|||||||
serde_yaml::to_string(&mixin).context("Failed to serialize generated sandbox mixin")
|
serde_yaml::to_string(&mixin).context("Failed to serialize generated sandbox mixin")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Renders the generated `coyote-mcp` mixin. Only proxy-managed credentials
|
||||||
|
/// are declared. sbx does not materialize env vars for `proxyManaged: false`
|
||||||
|
/// mixin credentials, so the rest are provisioned as custom secrets outside
|
||||||
|
/// the mixin, and only their target hosts join the network allow list here.
|
||||||
pub(crate) fn render_mixin_yaml(
|
pub(crate) fn render_mixin_yaml(
|
||||||
credentials: &[CredentialSpec],
|
credentials: &[CredentialSpec],
|
||||||
server_allow_entries: &[String],
|
server_allow_entries: &[String],
|
||||||
) -> Result<String> {
|
) -> Result<String> {
|
||||||
let entries = credentials
|
let entries = credentials
|
||||||
.iter()
|
.iter()
|
||||||
|
.filter(|c| c.proxy_managed)
|
||||||
.map(|c| CredentialEntry {
|
.map(|c| CredentialEntry {
|
||||||
service: c.service_id.clone(),
|
service: c.service_id.clone(),
|
||||||
description: format!(
|
description: format!(
|
||||||
@@ -510,14 +633,20 @@ pub(crate) fn render_mixin_yaml(
|
|||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
|
let mut allow_entries: Vec<String> = server_allow_entries.to_vec();
|
||||||
|
for credential in credentials.iter().filter(|c| !c.proxy_managed) {
|
||||||
|
allow_entries.extend(credential.custom_allow_entries.iter().cloned());
|
||||||
|
}
|
||||||
|
|
||||||
render_mixin_document(
|
render_mixin_document(
|
||||||
MCP_MIXIN_NAME,
|
MCP_MIXIN_NAME,
|
||||||
"Auto-generated by Coyote at launch: allows network egress to the user's remote MCP \
|
"Auto-generated by Coyote at launch: allows network egress to the user's remote MCP \
|
||||||
servers and declares their credentials so Docker Sandboxes binds them (bindings are \
|
servers and declares their credentials so Docker Sandboxes binds them (bindings are \
|
||||||
approved on first interactive run). Values are pre-seeded from Coyote's vault via \
|
approved on first interactive run). Proxy-injectable values are pre-seeded from \
|
||||||
`sbx secret set`.",
|
Coyote's vault via `sbx secret set`; the remaining secrets are provisioned as \
|
||||||
|
placeholder-based custom secrets via `sbx secret set-custom`.",
|
||||||
entries,
|
entries,
|
||||||
server_allow_entries,
|
&allow_entries,
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -603,6 +732,11 @@ mod tests {
|
|||||||
assert_eq!(cred.service_id, "github-pat");
|
assert_eq!(cred.service_id, "github-pat");
|
||||||
assert_eq!(cred.env_var, "COYOTE_SECRET_GITHUB_PAT");
|
assert_eq!(cred.env_var, "COYOTE_SECRET_GITHUB_PAT");
|
||||||
assert!(cred.proxy_managed);
|
assert!(cred.proxy_managed);
|
||||||
|
assert!(
|
||||||
|
cred.custom_hosts.is_empty(),
|
||||||
|
"proxy-managed credentials are provisioned via `sbx secret set`, \
|
||||||
|
not set-custom, so they carry no custom hosts"
|
||||||
|
);
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
cred.inject,
|
cred.inject,
|
||||||
vec![InjectRule {
|
vec![InjectRule {
|
||||||
@@ -839,6 +973,12 @@ mod tests {
|
|||||||
"secrets sharing an inject domain must both fall back to env"
|
"secrets sharing an inject domain must both fall back to env"
|
||||||
);
|
);
|
||||||
assert!(creds.iter().all(|c| c.inject.is_empty()));
|
assert!(creds.iter().all(|c| c.inject.is_empty()));
|
||||||
|
assert!(
|
||||||
|
creds
|
||||||
|
.iter()
|
||||||
|
.all(|c| c.custom_hosts == vec!["api.githubcopilot.com".to_string()]),
|
||||||
|
"demoted secrets must derive their set-custom targets from the server url"
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -1018,21 +1158,162 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn rendered_mixin_omits_inject_and_permissions_for_env_based_secrets() {
|
fn rendered_mixin_drops_env_based_credentials() {
|
||||||
let servers = servers(json!({
|
let servers = servers(json!({
|
||||||
"local": { "command": "run", "env": { "KEY": "{{NOTION_TOKEN}}" } }
|
"local": { "command": "run", "env": { "KEY": "{{NOTION_TOKEN}}" } }
|
||||||
}));
|
}));
|
||||||
|
|
||||||
let creds = collect_credentials(&servers).unwrap();
|
let creds = collect_credentials(&servers).unwrap();
|
||||||
|
assert_eq!(creds.len(), 1);
|
||||||
|
assert!(!creds[0].proxy_managed, "precondition");
|
||||||
|
assert!(
|
||||||
|
creds[0].custom_hosts.is_empty(),
|
||||||
|
"a stdio server with no URLs anywhere yields no derivable hosts"
|
||||||
|
);
|
||||||
|
|
||||||
let yaml = render_mixin_yaml(&creds, &[]).unwrap();
|
let yaml = render_mixin_yaml(&creds, &[]).unwrap();
|
||||||
let value: serde_yaml::Value = serde_yaml::from_str(&yaml).unwrap();
|
let value: serde_yaml::Value = serde_yaml::from_str(&yaml).unwrap();
|
||||||
|
|
||||||
let cred = &value["credentials"][0];
|
assert!(
|
||||||
assert_eq!(cred["apiKey"]["proxyManaged"].as_bool(), Some(false));
|
value.get("credentials").is_none(),
|
||||||
assert!(cred["apiKey"].get("inject").is_none());
|
"sbx does not materialize env vars for proxyManaged:false mixin \
|
||||||
|
credentials; declaring them would be dead config"
|
||||||
|
);
|
||||||
assert!(value.get("permissions").is_none());
|
assert!(value.get("permissions").is_none());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn hosts_in_text_extracts_hosts_and_strips_ports() {
|
||||||
|
assert_eq!(
|
||||||
|
hosts_in_text("https://api.example.com:8443/v1"),
|
||||||
|
BTreeSet::from(["api.example.com".to_string()])
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
hosts_in_text("see http://a.example.com/x and https://b.example.com/y"),
|
||||||
|
BTreeSet::from(["a.example.com".to_string(), "b.example.com".to_string()])
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
hosts_in_text("https://api.example.com/mcp?key={{KEY}}"),
|
||||||
|
BTreeSet::from(["api.example.com".to_string()])
|
||||||
|
);
|
||||||
|
assert!(hosts_in_text("ws://sock.example.com/mcp").is_empty());
|
||||||
|
assert!(hosts_in_text("qdrant.example.com:6333").is_empty());
|
||||||
|
assert!(hosts_in_text("https://[::1]:8443/mcp").is_empty());
|
||||||
|
assert!(hosts_in_text("httpserver is not a scheme").is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn hosts_in_text_terminates_urls_at_punctuation() {
|
||||||
|
assert_eq!(
|
||||||
|
hosts_in_text("https://a.example.com,https://b.example.com;https://c.example.com"),
|
||||||
|
BTreeSet::from([
|
||||||
|
"a.example.com".to_string(),
|
||||||
|
"b.example.com".to_string(),
|
||||||
|
"c.example.com".to_string()
|
||||||
|
])
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
hosts_in_text("(see https://docs.example.com)"),
|
||||||
|
BTreeSet::from(["docs.example.com".to_string()])
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn server_secret_hosts_unions_url_host_and_scraped_urls() {
|
||||||
|
let server = json!({
|
||||||
|
"command": "run",
|
||||||
|
"args": ["--endpoint", "https://api.vendor.example/v2"],
|
||||||
|
"env": { "BASE_URL": "http://internal.example.com:8080/api" }
|
||||||
|
});
|
||||||
|
|
||||||
|
let hosts = server_secret_hosts(&server);
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
hosts.targets,
|
||||||
|
BTreeSet::from([
|
||||||
|
"api.vendor.example".to_string(),
|
||||||
|
"internal.example.com".to_string()
|
||||||
|
]),
|
||||||
|
"set-custom targets are port-stripped"
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
hosts.allow_entries,
|
||||||
|
BTreeSet::from([
|
||||||
|
"api.vendor.example".to_string(),
|
||||||
|
"internal.example.com:8080".to_string()
|
||||||
|
]),
|
||||||
|
"allow entries keep non-default ports so the hosts stay reachable"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn server_secret_hosts_includes_non_http_url_host() {
|
||||||
|
let server = json!({ "url": "ws://sock.example.com/mcp" });
|
||||||
|
|
||||||
|
let hosts = server_secret_hosts(&server);
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
hosts.targets,
|
||||||
|
BTreeSet::from(["sock.example.com".to_string()])
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
hosts.allow_entries.is_empty(),
|
||||||
|
"non-http(s) urls have no spelling in the allow grammar"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn server_secret_hosts_skips_placeholder_hosts() {
|
||||||
|
let server = json!({
|
||||||
|
"url": "https://{{TENANT}}.example.com/mcp",
|
||||||
|
"env": { "API_URL": "https://{{REGION}}.api.example.com/v1" }
|
||||||
|
});
|
||||||
|
|
||||||
|
let hosts = server_secret_hosts(&server);
|
||||||
|
|
||||||
|
assert!(
|
||||||
|
hosts.targets.is_empty() && hosts.allow_entries.is_empty(),
|
||||||
|
"a host containing an unresolved placeholder must never reach \
|
||||||
|
set-custom argv or the mixin allow list: {hosts:?}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn demoted_credential_hosts_are_unioned_into_allow() {
|
||||||
|
let servers = servers(json!({
|
||||||
|
"local": {
|
||||||
|
"command": "run",
|
||||||
|
"env": {
|
||||||
|
"KEY": "{{TOKEN}}",
|
||||||
|
"API_URL": "https://api.internal.example.com:8443/v1"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}));
|
||||||
|
|
||||||
|
let creds = collect_credentials(&servers).unwrap();
|
||||||
|
assert_eq!(creds.len(), 1);
|
||||||
|
assert!(!creds[0].proxy_managed, "precondition");
|
||||||
|
assert_eq!(
|
||||||
|
creds[0].custom_hosts,
|
||||||
|
vec!["api.internal.example.com".to_string()]
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
creds[0].custom_allow_entries,
|
||||||
|
vec!["api.internal.example.com:8443".to_string()],
|
||||||
|
"allow entries keep the port that set-custom targets must strip"
|
||||||
|
);
|
||||||
|
|
||||||
|
let yaml = render_mixin_yaml(&creds, &[]).unwrap();
|
||||||
|
let value: serde_yaml::Value = serde_yaml::from_str(&yaml).unwrap();
|
||||||
|
|
||||||
|
assert!(value.get("credentials").is_none());
|
||||||
|
assert_eq!(
|
||||||
|
value["permissions"]["network"]["allow"][0].as_str(),
|
||||||
|
Some("api.internal.example.com:8443"),
|
||||||
|
"a demoted credential's derived hosts must still be reachable"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn rendered_mixin_is_deterministic() {
|
fn rendered_mixin_is_deterministic() {
|
||||||
let servers = servers(json!({
|
let servers = servers(json!({
|
||||||
|
|||||||
+816
-73
File diff suppressed because it is too large
Load Diff
+29
-1
@@ -35,6 +35,7 @@ use nu_ansi_term::Color;
|
|||||||
use serde_json::Value;
|
use serde_json::Value;
|
||||||
use std::borrow::Cow;
|
use std::borrow::Cow;
|
||||||
use std::collections::VecDeque;
|
use std::collections::VecDeque;
|
||||||
|
use std::io;
|
||||||
use std::sync::atomic::AtomicBool;
|
use std::sync::atomic::AtomicBool;
|
||||||
use std::sync::{LazyLock, Mutex, OnceLock};
|
use std::sync::{LazyLock, Mutex, OnceLock};
|
||||||
use std::{cmp, env, path::PathBuf, process};
|
use std::{cmp, env, path::PathBuf, process};
|
||||||
@@ -45,7 +46,7 @@ pub static CODE_BLOCK_RE: LazyLock<Regex> =
|
|||||||
LazyLock::new(|| Regex::new(r"(?ms)```\w*(.*)```").unwrap());
|
LazyLock::new(|| Regex::new(r"(?ms)```\w*(.*)```").unwrap());
|
||||||
pub static THINK_TAG_RE: LazyLock<Regex> =
|
pub static THINK_TAG_RE: LazyLock<Regex> =
|
||||||
LazyLock::new(|| Regex::new(r"(?s)^\s*<think>.*?</think>(\s*|$)").unwrap());
|
LazyLock::new(|| Regex::new(r"(?s)^\s*<think>.*?</think>(\s*|$)").unwrap());
|
||||||
pub static IS_STDOUT_TERMINAL: LazyLock<bool> = LazyLock::new(|| std::io::stdout().is_terminal());
|
pub static IS_STDOUT_TERMINAL: LazyLock<bool> = LazyLock::new(|| io::stdout().is_terminal());
|
||||||
pub static HEADLESS: AtomicBool = AtomicBool::new(false);
|
pub static HEADLESS: AtomicBool = AtomicBool::new(false);
|
||||||
pub static ACP_SERVER: AtomicBool = AtomicBool::new(false);
|
pub static ACP_SERVER: AtomicBool = AtomicBool::new(false);
|
||||||
|
|
||||||
@@ -133,6 +134,33 @@ pub fn parse_bool(value: &str) -> Option<bool> {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn drain_stale_tty_input() {
|
||||||
|
use crossterm::event::{poll, read};
|
||||||
|
use std::time::{Duration, Instant};
|
||||||
|
|
||||||
|
if !io::stdin().is_terminal() {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
if crossterm::terminal::enable_raw_mode().is_err() {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
let deadline = Instant::now() + Duration::from_millis(100);
|
||||||
|
while Instant::now() < deadline {
|
||||||
|
match poll(Duration::from_millis(10)) {
|
||||||
|
Ok(true) => {
|
||||||
|
if read().is_err() {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => break,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let _ = crossterm::terminal::disable_raw_mode();
|
||||||
|
}
|
||||||
|
|
||||||
pub fn estimate_token_length(text: &str) -> usize {
|
pub fn estimate_token_length(text: &str) -> usize {
|
||||||
let weighted: usize = text.chars().map(|c| if c.is_ascii() { 1 } else { 2 }).sum();
|
let weighted: usize = text.chars().map(|c| if c.is_ascii() { 1 } else { 2 }).sum();
|
||||||
weighted.div_ceil(4)
|
weighted.div_ceil(4)
|
||||||
|
|||||||
@@ -54,20 +54,6 @@ pub async fn expand_glob_paths<T: AsRef<str>>(
|
|||||||
Ok(new_paths)
|
Ok(new_paths)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn clear_dir(dir: &Path) -> Result<()> {
|
|
||||||
for entry in fs::read_dir(dir)? {
|
|
||||||
let entry = entry?;
|
|
||||||
let path = entry.path();
|
|
||||||
|
|
||||||
if path.is_dir() {
|
|
||||||
fs::remove_dir_all(&path)?;
|
|
||||||
} else {
|
|
||||||
fs::remove_file(&path)?;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn list_file_names<T: AsRef<Path>>(dir: T, ext: &str) -> Vec<String> {
|
pub fn list_file_names<T: AsRef<Path>>(dir: T, ext: &str) -> Vec<String> {
|
||||||
match fs::read_dir(dir.as_ref()) {
|
match fs::read_dir(dir.as_ref()) {
|
||||||
Ok(rd) => {
|
Ok(rd) => {
|
||||||
|
|||||||
@@ -12,6 +12,7 @@ pub use utils::prompt_provider_choice;
|
|||||||
|
|
||||||
use crate::cli::Cli;
|
use crate::cli::Cli;
|
||||||
use crate::config::AppConfig;
|
use crate::config::AppConfig;
|
||||||
|
use crate::utils::drain_stale_tty_input;
|
||||||
use crate::vault::utils::ensure_password_file_initialized;
|
use crate::vault::utils::ensure_password_file_initialized;
|
||||||
use anyhow::{Context, Result, anyhow, bail};
|
use anyhow::{Context, Result, anyhow, bail};
|
||||||
use fancy_regex::Regex;
|
use fancy_regex::Regex;
|
||||||
@@ -151,6 +152,7 @@ impl Vault {
|
|||||||
"Vault management is disabled in sandbox mode. Use `coyote --add-secret` on your host."
|
"Vault management is disabled in sandbox mode. Use `coyote --add-secret` on your host."
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
drain_stale_tty_input();
|
||||||
let secret_value = Password::new("Enter the secret value:")
|
let secret_value = Password::new("Enter the secret value:")
|
||||||
.with_validator(required!())
|
.with_validator(required!())
|
||||||
.with_display_mode(PasswordDisplayMode::Masked)
|
.with_display_mode(PasswordDisplayMode::Masked)
|
||||||
@@ -190,6 +192,7 @@ impl Vault {
|
|||||||
"Vault management is disabled in sandbox mode. Use `coyote --add-secret` on your host."
|
"Vault management is disabled in sandbox mode. Use `coyote --add-secret` on your host."
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
drain_stale_tty_input();
|
||||||
let secret_value = Password::new("Enter the secret value:")
|
let secret_value = Password::new("Enter the secret value:")
|
||||||
.with_validator(required!())
|
.with_validator(required!())
|
||||||
.with_display_mode(PasswordDisplayMode::Masked)
|
.with_display_mode(PasswordDisplayMode::Masked)
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
use crate::config::ensure_parent_exists;
|
use crate::config::ensure_parent_exists;
|
||||||
use crate::sandbox::{SANDBOX_ENV_FLAG, sandbox_secret_env_var};
|
use crate::sandbox::{SANDBOX_ENV_FLAG, sandbox_secret_env_var};
|
||||||
|
use crate::utils::drain_stale_tty_input;
|
||||||
use crate::vault::{SECRET_RE, Vault};
|
use crate::vault::{SECRET_RE, Vault};
|
||||||
use anyhow::Result;
|
use anyhow::Result;
|
||||||
use anyhow::anyhow;
|
use anyhow::anyhow;
|
||||||
@@ -68,6 +69,7 @@ pub fn create_vault_password_file(vault: &mut Vault) -> Result<()> {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
drain_stale_tty_input();
|
||||||
let ans = Confirm::new(
|
let ans = Confirm::new(
|
||||||
format!(
|
format!(
|
||||||
"The configured password file '{}' is empty. Create a password?",
|
"The configured password file '{}' is empty. Create a password?",
|
||||||
@@ -107,6 +109,7 @@ pub fn create_vault_password_file(vault: &mut Vault) -> Result<()> {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
|
drain_stale_tty_input();
|
||||||
let ans = Confirm::new("No password file configured. Do you want to create one now?")
|
let ans = Confirm::new("No password file configured. Do you want to create one now?")
|
||||||
.with_default(true)
|
.with_default(true)
|
||||||
.prompt()?;
|
.prompt()?;
|
||||||
@@ -185,6 +188,7 @@ pub fn create_vault_password_file(vault: &mut Vault) -> Result<()> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub fn prompt_provider_choice() -> Result<Option<SupportedProvider>> {
|
pub fn prompt_provider_choice() -> Result<Option<SupportedProvider>> {
|
||||||
|
drain_stale_tty_input();
|
||||||
let choices = vec![
|
let choices = vec![
|
||||||
"local - encrypted file on this machine",
|
"local - encrypted file on this machine",
|
||||||
"aws_secrets_manager - AWS Secrets Manager",
|
"aws_secrets_manager - AWS Secrets Manager",
|
||||||
|
|||||||
Reference in New Issue
Block a user