review classifier: run on ACP when the swarm settings choose it
The swarm settings gain `runtime` (`api-key`, the default, or `acp`) and an `acp` provider record: an https base URL, the provider's model, and its prices per million tokens. A PUT now changes only the fields its body names, so the settings form that sends model, effort and limit leaves the runtime alone; `null` clears `effort` or `acp`. The body stays strict, and the stored document keeps #5192's lenient read. The controller sends the runtime in each run request. On `acp` the classifier runs one prompt on a new `opencode acp` session per classification through hive-runtime's AcpRuntime, in a scratch directory, with a config rendered from the settings. opencode runs `read` and other tools without asking, so OPENCODE_PERMISSION {"*":"deny"} takes every tool away from the model, and a deny-all policy refuses any request on top. The title call, which resends the prompt, is turned off. The deadline cancels the turn, and a turn that does not end within 15s of that is dropped, which also covers an agent stuck before its prompt. The answer is the final assistant text, as ACP has no result event. The provider key is read per run from swarm/services/review-classifier/acp-provider, which the bao UI's admins login may create and update, and reaches opencode through AcpRuntime::with_provider_key and the trampoline's allowlist. Each ACP classification's cost is computed from the configured prices and the reported tokens, and added to hyperhive.review_classifier.cost.usage labelled repo, pr, model and currency. Part of #5096
This commit is contained in:
parent
38ee68c31d
commit
1b4e764e5c
24 changed files with 1622 additions and 262 deletions
|
|
@ -13,9 +13,15 @@ anyhow.workspace = true
|
|||
async-nats.workspace = true
|
||||
# StreamExt::next on the subscription.
|
||||
futures-util.workspace = true
|
||||
# One `claude --print` per classification.
|
||||
# One `claude --print` per classification on the API key runtime.
|
||||
hive-claude.workspace = true
|
||||
hive-log.workspace = true
|
||||
# One ACP session per classification on the ACP runtime.
|
||||
hive-runtime.workspace = true
|
||||
# The ACP runtime's cost counter (`cost`).
|
||||
opentelemetry.workspace = true
|
||||
opentelemetry_sdk.workspace = true
|
||||
opentelemetry-otlp.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
# The try-out subject and the shapes the controller sends and reads back.
|
||||
|
|
@ -25,5 +31,8 @@ swarm-secret-client.workspace = true
|
|||
tokio.workspace = true
|
||||
tracing.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = "3"
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
|
|
|||
467
swarm-review-classifier/src/acp.rs
Normal file
467
swarm-review-classifier/src/acp.rs
Normal file
|
|
@ -0,0 +1,467 @@
|
|||
//! One classification on the ACP runtime: `opencode acp`, driven through
|
||||
//! [`hive_runtime::AcpRuntime`] for one prompt on a new session, in a
|
||||
//! directory made for it and removed after it.
|
||||
//!
|
||||
//! The agent has no tool at all. opencode runs some tools, `read` among them,
|
||||
//! without asking, so a deny-all permission policy alone would never see them:
|
||||
//! [`DENY_ALL`] in `OPENCODE_PERMISSION` makes opencode offer the model no
|
||||
//! tool, and the policy refuses whatever is still asked. A tool would let a
|
||||
//! PR's diff talk the model into reading this unit's credentials into its
|
||||
//! `reason`, which the controller posts on the forge.
|
||||
//!
|
||||
//! The rendered config names one provider and model, from the run, and turns
|
||||
//! off opencode's title call, which would send the prompt a second time.
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
|
||||
use hive_runtime::{
|
||||
AcpCommand, AcpRuntime, Config, PermissionAsk, PermissionPolicy, Runtime as _, Sink,
|
||||
};
|
||||
use serde_json::{Value, json};
|
||||
use swarm_queue_client::review_classifier::{Acp as Provider, Runtime, TryOut};
|
||||
|
||||
use crate::cost::{Cost, CostSink};
|
||||
|
||||
/// The variable the provider key reaches `opencode` in.
|
||||
pub const PROVIDER_KEY_VAR: &str = "ACP_PROVIDER_API_KEY";
|
||||
|
||||
/// `OPENCODE_PERMISSION`: every tool denied, which opencode merges over its
|
||||
/// config.
|
||||
const DENY_ALL: &str = r#"{"*":"deny"}"#;
|
||||
|
||||
/// The provider's id in the rendered config.
|
||||
const PROVIDER_ID: &str = "review-classifier";
|
||||
|
||||
/// How long the agent may send nothing before its turn is cancelled.
|
||||
const IDLE_TIMEOUT: Duration = Duration::from_mins(2);
|
||||
|
||||
/// How long a turn cancelled at the deadline has to end before it is dropped,
|
||||
/// which interrupts the agent: past `hive-runtime`'s own grace before it sends
|
||||
/// SIGINT. A cancel reaches only a turn whose prompt has been sent, so an
|
||||
/// agent stuck before that needs the drop.
|
||||
pub const CANCEL_WAIT: Duration = Duration::from_secs(15);
|
||||
|
||||
/// The `opencode` binary the classifications run on, started through the
|
||||
/// trampoline that clears its environment.
|
||||
pub struct Acp<S> {
|
||||
pub program: PathBuf,
|
||||
/// This binary, which execs [`Self::program`] when started as the
|
||||
/// trampoline.
|
||||
pub trampoline: PathBuf,
|
||||
/// How long one classification may take before it is cancelled.
|
||||
pub deadline: Duration,
|
||||
/// [`CANCEL_WAIT`] outside the tests.
|
||||
pub cancel_wait: Duration,
|
||||
pub cost: S,
|
||||
}
|
||||
|
||||
impl<S: CostSink> crate::run::Model for Acp<S> {
|
||||
async fn ask(&self, prompt: &str, key: &str, run: &TryOut) -> Result<String, String> {
|
||||
let Runtime::Acp(provider) = &run.runtime else {
|
||||
return Err("not an ACP run".to_owned());
|
||||
};
|
||||
let dir = crate::run::scratch_dir();
|
||||
let outcome = self.ask_in(&dir, prompt, key, run, provider).await;
|
||||
if let Err(e) = tokio::fs::remove_dir_all(&dir).await {
|
||||
tracing::warn!(dir = %dir.display(), error = %e, "removing a run's directory failed");
|
||||
}
|
||||
outcome
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: CostSink> Acp<S> {
|
||||
async fn ask_in(
|
||||
&self,
|
||||
dir: &Path,
|
||||
prompt: &str,
|
||||
key: &str,
|
||||
run: &TryOut,
|
||||
provider: &Provider,
|
||||
) -> Result<String, String> {
|
||||
let system = dir.join("system.md");
|
||||
let config_file = dir.join("opencode.json");
|
||||
tokio::fs::create_dir_all(dir)
|
||||
.await
|
||||
.and(tokio::fs::write(&system, crate::prompt::SYSTEM).await)
|
||||
.and(tokio::fs::write(&config_file, opencode_config(provider).to_string()).await)
|
||||
.map_err(|e| format!("preparing the run's directory: {e}"))?;
|
||||
let runtime = AcpRuntime::new(
|
||||
self.command(dir, &config_file),
|
||||
dir.join("session"),
|
||||
deny_all(),
|
||||
hive_claude::NeverCompact,
|
||||
)
|
||||
.with_provider_key(key.to_owned());
|
||||
let canceller = runtime
|
||||
.canceller()
|
||||
.ok_or("the ACP runtime offers no cancel handle")?;
|
||||
let config = Config {
|
||||
effort: run.effort.map(|e| e.as_str().to_owned()),
|
||||
cwd: Some(dir.to_owned()),
|
||||
system_prompt_file: Some(system),
|
||||
idle_timeout: Some(IDLE_TIMEOUT),
|
||||
..Config::default()
|
||||
};
|
||||
let sink = FinalText::default();
|
||||
let turn = runtime.run(&config, prompt, &sink);
|
||||
tokio::pin!(turn);
|
||||
match tokio::time::timeout(self.deadline, &mut turn).await {
|
||||
Ok(Ok(progress)) => {
|
||||
self.record_cost(run, provider, &progress.telemetry.cost);
|
||||
sink.take()
|
||||
}
|
||||
Ok(Err(e)) => Err(format!("the ACP agent failed: {e}")),
|
||||
Err(_) => {
|
||||
let _ = canceller.cancel();
|
||||
let _ = tokio::time::timeout(self.cancel_wait, turn).await;
|
||||
Err(format!(
|
||||
"no answer within {} seconds",
|
||||
self.deadline.as_secs()
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// One classification's invocation: the trampoline as the program, with
|
||||
/// the real `opencode` and the run's own variables in its environment.
|
||||
fn command(&self, dir: &Path, config_file: &Path) -> AcpCommand {
|
||||
let env: BTreeMap<String, String> = [
|
||||
(
|
||||
crate::child_env::EXEC_OPENCODE,
|
||||
self.program.display().to_string(),
|
||||
),
|
||||
("HOME", dir.display().to_string()),
|
||||
("OPENCODE_CONFIG", config_file.display().to_string()),
|
||||
("OPENCODE_PERMISSION", DENY_ALL.to_owned()),
|
||||
("OPENCODE_DISABLE_PROJECT_CONFIG", "1".to_owned()),
|
||||
]
|
||||
.into_iter()
|
||||
.map(|(k, v)| (k.to_owned(), v))
|
||||
.collect();
|
||||
AcpCommand {
|
||||
command: self.trampoline.display().to_string(),
|
||||
args: vec!["acp".to_owned()],
|
||||
env,
|
||||
api_key_env: Some(PROVIDER_KEY_VAR.to_owned()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Record what `usage` cost at the run's prices. An agent that reports no
|
||||
/// tokens leaves the classification's cost unknown, not zero.
|
||||
fn record_cost(&self, run: &TryOut, provider: &Provider, usage: &hive_runtime::TokenUsage) {
|
||||
if *usage == hive_runtime::TokenUsage::default() {
|
||||
tracing::warn!(
|
||||
repo = %run.repo, pr = run.pr,
|
||||
"the ACP agent reported no token counts, so this classification's cost is unknown"
|
||||
);
|
||||
return;
|
||||
}
|
||||
self.cost.record(&Cost {
|
||||
repo: run.repo.clone(),
|
||||
pr: run.pr,
|
||||
model: provider.model.clone(),
|
||||
amount: crate::cost::amount(&provider.prices, usage),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/// Refuses every permission request.
|
||||
fn deny_all() -> PermissionPolicy {
|
||||
Arc::new(|_: &PermissionAsk<'_>| false)
|
||||
}
|
||||
|
||||
/// The `opencode.json` of a run on `provider`.
|
||||
fn opencode_config(provider: &Provider) -> Value {
|
||||
let prices = &provider.prices;
|
||||
json!({
|
||||
"$schema": "https://opencode.ai/config.json",
|
||||
"autoupdate": false,
|
||||
"share": "disabled",
|
||||
"model": format!("{PROVIDER_ID}/{}", provider.model),
|
||||
"enabled_providers": [PROVIDER_ID],
|
||||
"agent": { "title": { "disable": true } },
|
||||
"compaction": { "auto": false },
|
||||
"provider": {
|
||||
PROVIDER_ID: {
|
||||
"npm": "@ai-sdk/openai-compatible",
|
||||
"name": "review classifier provider",
|
||||
"options": {
|
||||
"baseURL": provider.base_url,
|
||||
// opencode substitutes it from its environment.
|
||||
"apiKey": format!("{{env:{PROVIDER_KEY_VAR}}}"),
|
||||
},
|
||||
"models": {
|
||||
provider.model.as_str(): {
|
||||
"name": provider.model,
|
||||
"cost": {
|
||||
"input": prices.input,
|
||||
"output": prices.output,
|
||||
"cache_read": prices.cache_read.unwrap_or(prices.input),
|
||||
"cache_write": prices.cache_write.unwrap_or(prices.input),
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
/// Keeps the turn's last assistant text: the model's final message. An ACP
|
||||
/// turn has no `result` event to take it from.
|
||||
#[derive(Default)]
|
||||
struct FinalText(Mutex<Option<String>>);
|
||||
|
||||
impl FinalText {
|
||||
fn take(&self) -> Result<String, String> {
|
||||
self.0
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.take()
|
||||
.ok_or_else(|| "the ACP agent ended without an answer".to_owned())
|
||||
}
|
||||
}
|
||||
|
||||
impl Sink for FinalText {
|
||||
fn on_event(&self, event: &Value) {
|
||||
if event["type"] != "assistant" {
|
||||
return;
|
||||
}
|
||||
let mut texts = event["message"]["content"]
|
||||
.as_array()
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter(|block| block["type"] == "text")
|
||||
.filter_map(|block| block["text"].as_str());
|
||||
if let Some(text) = texts.next_back() {
|
||||
*self
|
||||
.0
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner) = Some(text.to_owned());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::os::unix::fs::PermissionsExt as _;
|
||||
|
||||
use swarm_queue_client::review_classifier::Prices;
|
||||
|
||||
use super::*;
|
||||
use crate::run::Model as _;
|
||||
|
||||
/// An ACP agent in plain `sh`, logging to `<log>.*`. It writes its
|
||||
/// environment and a copy of its config at start, appends each method
|
||||
/// it is sent to `<log>.methods`, and answers a prompt with a text chunk
|
||||
/// holding a verdict, then `end_turn` with 1000 input and 200 output
|
||||
/// tokens. A prompt holding `ASK-READ` first asks permission for a
|
||||
/// `read` and appends the reply to `<log>.permission`; one holding
|
||||
/// `STALL` is answered only once it is cancelled; one holding `MUTE`
|
||||
/// gets the answer with no token counts.
|
||||
const AGENT: &str = r#"#!/bin/sh
|
||||
log="$1"
|
||||
env > "$log.env"
|
||||
cp "$OPENCODE_CONFIG" "$log.config"
|
||||
stalled=
|
||||
while IFS= read -r line; do
|
||||
id=${line#*\"id\":}; id=${id%%[,\}]*}
|
||||
method=$(printf '%s' "$line" | sed -n 's/.*"method":"\([^"]*\)".*/\1/p')
|
||||
printf '%s\n' "$method" >> "$log.methods"
|
||||
case "$method" in
|
||||
initialize)
|
||||
printf '{"jsonrpc":"2.0","id":%s,"result":{"protocolVersion":1,"agentCapabilities":{"mcpCapabilities":{"http":true}}}}\n' "$id" ;;
|
||||
session/new)
|
||||
printf '{"jsonrpc":"2.0","id":%s,"result":{"sessionId":"s1"}}\n' "$id" ;;
|
||||
session/cancel)
|
||||
[ -n "$stalled" ] && printf '{"jsonrpc":"2.0","id":%s,"result":{"stopReason":"cancelled"}}\n' "$stalled" ;;
|
||||
session/prompt)
|
||||
case "$line" in *STALL*) stalled=$id; continue ;; esac
|
||||
case "$line" in *ASK-READ*)
|
||||
printf '{"jsonrpc":"2.0","id":900,"method":"session/request_permission","params":{"sessionId":"s1","toolCall":{"toolCallId":"t1","kind":"read","title":"read"},"options":[{"optionId":"yes","kind":"allow_once","name":"Allow"},{"optionId":"no","kind":"reject_once","name":"Reject"}]}}\n'
|
||||
read -r reply
|
||||
printf '%s\n' "$reply" >> "$log.permission" ;;
|
||||
esac
|
||||
printf '{"jsonrpc":"2.0","method":"session/update","params":{"sessionId":"s1","update":{"sessionUpdate":"agent_message_chunk","content":{"type":"text","text":"{\\"verdict\\":\\"match\\",\\"certainty\\":90,\\"reason\\":\\"ok\\"}"}}}}\n'
|
||||
case "$line" in
|
||||
*MUTE*) printf '{"jsonrpc":"2.0","id":%s,"result":{"stopReason":"end_turn"}}\n' "$id" ;;
|
||||
*) printf '{"jsonrpc":"2.0","id":%s,"result":{"stopReason":"end_turn","usage":{"inputTokens":1000,"outputTokens":200}}}\n' "$id" ;;
|
||||
esac ;;
|
||||
esac
|
||||
done
|
||||
"#;
|
||||
|
||||
#[derive(Default)]
|
||||
struct Costs(Mutex<Vec<Cost>>);
|
||||
|
||||
impl CostSink for Arc<Costs> {
|
||||
fn record(&self, cost: &Cost) {
|
||||
self.0.lock().expect("unpoisoned").push(cost.clone());
|
||||
}
|
||||
}
|
||||
|
||||
/// An [`Acp`] whose trampoline is [`AGENT`], logging under `dir`.
|
||||
fn acp(dir: &Path, deadline: Duration) -> (Acp<Arc<Costs>>, Arc<Costs>) {
|
||||
let script = dir.join("agent.sh");
|
||||
let log = dir.join("log");
|
||||
std::fs::write(
|
||||
&script,
|
||||
AGENT.replacen("log=\"$1\"", &format!("log='{}'", log.display()), 1),
|
||||
)
|
||||
.expect("written");
|
||||
std::fs::set_permissions(&script, std::fs::Permissions::from_mode(0o755)).expect("chmod");
|
||||
let costs = Arc::new(Costs::default());
|
||||
let acp = Acp {
|
||||
program: PathBuf::from("/nix/store/x-opencode/bin/opencode"),
|
||||
trampoline: script,
|
||||
deadline,
|
||||
cancel_wait: Duration::from_millis(500),
|
||||
cost: Arc::clone(&costs),
|
||||
};
|
||||
(acp, costs)
|
||||
}
|
||||
|
||||
fn read(dir: &Path, suffix: &str) -> String {
|
||||
std::fs::read_to_string(dir.join(format!("log.{suffix}"))).unwrap_or_default()
|
||||
}
|
||||
|
||||
fn run() -> TryOut {
|
||||
TryOut {
|
||||
operator: "alice".to_owned(),
|
||||
repo: "o/r".to_owned(),
|
||||
pr: 7,
|
||||
head_sha: "abc".to_owned(),
|
||||
commits: Vec::new(),
|
||||
diff: String::new(),
|
||||
classifications: Vec::new(),
|
||||
model: "haiku".to_owned(),
|
||||
effort: None,
|
||||
runtime: Runtime::Acp(Provider {
|
||||
base_url: "https://llm.example/v1".to_owned(),
|
||||
model: "qwen".to_owned(),
|
||||
prices: Prices {
|
||||
input: 3.0,
|
||||
output: 15.0,
|
||||
cache_read: None,
|
||||
cache_write: None,
|
||||
},
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn the_answer_is_the_final_text_and_its_cost_is_labelled_with_the_pr() {
|
||||
let dir = tempfile::tempdir().expect("a dir");
|
||||
let (acp, costs) = acp(dir.path(), Duration::from_mins(1));
|
||||
|
||||
let answer = acp.ask("classify", "p", &run()).await;
|
||||
assert_eq!(
|
||||
answer.as_deref(),
|
||||
Ok(r#"{"verdict":"match","certainty":90,"reason":"ok"}"#)
|
||||
);
|
||||
let costs = costs.0.lock().expect("unpoisoned");
|
||||
assert_eq!(costs.len(), 1);
|
||||
let cost = &costs[0];
|
||||
assert_eq!(
|
||||
(cost.repo.as_str(), cost.pr, cost.model.as_str()),
|
||||
("o/r", 7, "qwen")
|
||||
);
|
||||
assert!((cost.amount - 0.006).abs() < 1e-12, "{}", cost.amount);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn opencode_gets_the_key_no_tools_and_the_runs_provider() {
|
||||
let dir = tempfile::tempdir().expect("a dir");
|
||||
let (acp, _) = acp(dir.path(), Duration::from_mins(1));
|
||||
|
||||
acp.ask("classify", "p", &run()).await.expect("an answer");
|
||||
let env = read(dir.path(), "env");
|
||||
for line in [
|
||||
"ACP_PROVIDER_API_KEY=p",
|
||||
r#"OPENCODE_PERMISSION={"*":"deny"}"#,
|
||||
"OPENCODE_DISABLE_PROJECT_CONFIG=1",
|
||||
"SWARM_REVIEW_CLASSIFIER_EXEC_OPENCODE=/nix/store/x-opencode/bin/opencode",
|
||||
] {
|
||||
assert!(env.lines().any(|l| l == line), "{line} not in:\n{env}");
|
||||
}
|
||||
let config: Value = serde_json::from_str(&read(dir.path(), "config")).expect("JSON");
|
||||
let provider = &config["provider"][PROVIDER_ID];
|
||||
assert_eq!(config["model"], "review-classifier/qwen");
|
||||
assert_eq!(provider["options"]["baseURL"], "https://llm.example/v1");
|
||||
assert_eq!(provider["options"]["apiKey"], "{env:ACP_PROVIDER_API_KEY}");
|
||||
assert_eq!(provider["models"]["qwen"]["cost"]["output"], 15.0);
|
||||
assert_eq!(config["agent"]["title"]["disable"], true);
|
||||
// The run's directory is gone with the run.
|
||||
let home = env
|
||||
.lines()
|
||||
.find_map(|l| l.strip_prefix("HOME="))
|
||||
.expect("a HOME");
|
||||
assert!(!Path::new(home).exists(), "{home}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_permission_request_is_refused() {
|
||||
let dir = tempfile::tempdir().expect("a dir");
|
||||
let (acp, _) = acp(dir.path(), Duration::from_mins(1));
|
||||
|
||||
acp.ask("ASK-READ", "p", &run()).await.expect("an answer");
|
||||
let reply: Value =
|
||||
serde_json::from_str(read(dir.path(), "permission").trim()).expect("a JSON reply");
|
||||
assert_eq!(reply["result"]["outcome"]["optionId"], "no");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn the_deadline_cancels_the_turn() {
|
||||
let dir = tempfile::tempdir().expect("a dir");
|
||||
let (acp, costs) = acp(dir.path(), Duration::from_millis(500));
|
||||
|
||||
let answer = acp.ask("STALL", "p", &run()).await;
|
||||
assert_eq!(answer, Err("no answer within 0 seconds".to_owned()));
|
||||
assert!(
|
||||
read(dir.path(), "methods")
|
||||
.lines()
|
||||
.any(|m| m == "session/cancel")
|
||||
);
|
||||
assert!(costs.0.lock().expect("unpoisoned").is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn an_agent_stuck_before_its_prompt_is_dropped_at_the_deadline() {
|
||||
let dir = tempfile::tempdir().expect("a dir");
|
||||
let (acp, _) = acp(dir.path(), Duration::from_millis(500));
|
||||
let script = dir.path().join("agent.sh");
|
||||
let unanswered = std::fs::read_to_string(&script).expect("written").replacen(
|
||||
" initialize)",
|
||||
" never)",
|
||||
1,
|
||||
);
|
||||
std::fs::write(&script, unanswered).expect("written");
|
||||
|
||||
let answer = tokio::time::timeout(Duration::from_secs(10), acp.ask("one", "p", &run()))
|
||||
.await
|
||||
.expect("the deadline ends the ask");
|
||||
assert_eq!(answer, Err("no answer within 0 seconds".to_owned()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn no_token_counts_record_no_cost() {
|
||||
let dir = tempfile::tempdir().expect("a dir");
|
||||
let (acp, costs) = acp(dir.path(), Duration::from_mins(1));
|
||||
|
||||
acp.ask("MUTE", "p", &run()).await.expect("an answer");
|
||||
assert!(costs.0.lock().expect("unpoisoned").is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_claude_run_is_refused() {
|
||||
let dir = tempfile::tempdir().expect("a dir");
|
||||
let (acp, _) = acp(dir.path(), Duration::from_mins(1));
|
||||
let mut claude = run();
|
||||
claude.runtime = Runtime::ApiKey;
|
||||
|
||||
assert!(acp.ask("classify", "p", &claude).await.is_err());
|
||||
assert_eq!(read(dir.path(), "methods"), "", "no agent started");
|
||||
}
|
||||
}
|
||||
|
|
@ -1,10 +1,11 @@
|
|||
//! The `claude` child's environment: an allowlist, not the daemon's.
|
||||
//! The model child's environment: an allowlist, not the daemon's.
|
||||
//!
|
||||
//! `hive-claude` sets its variables on top of the parent's environment, and
|
||||
//! the daemon's holds the store identity paths and the queue coordinates. So
|
||||
//! `claude` is spawned as this same binary with [`EXEC_CLAUDE`] naming the
|
||||
//! real one: [`exec_if_requested`] runs first in `main`, clears the
|
||||
//! environment down to [`ALLOWED`] and execs `claude` with the same
|
||||
//! `hive-claude` and `hive-runtime` set their variables on top of the
|
||||
//! parent's environment, and the daemon's holds the store identity paths and
|
||||
//! the queue coordinates. So the child is spawned as this same binary with
|
||||
//! [`EXEC_CLAUDE`] or [`EXEC_OPENCODE`] naming the real one:
|
||||
//! [`exec_if_requested`] runs first in `main`, clears the environment down to
|
||||
//! that runtime's allowlist and execs the real binary with the same
|
||||
//! arguments, in the same process.
|
||||
|
||||
use std::ffi::OsString;
|
||||
|
|
@ -14,9 +15,12 @@ use std::process::Command;
|
|||
/// Set on the child only: the `claude` binary to exec.
|
||||
pub const EXEC_CLAUDE: &str = "SWARM_REVIEW_CLASSIFIER_EXEC_CLAUDE";
|
||||
|
||||
/// Set on the child only: the `opencode` binary to exec.
|
||||
pub const EXEC_OPENCODE: &str = "SWARM_REVIEW_CLASSIFIER_EXEC_OPENCODE";
|
||||
|
||||
/// Every variable `claude` keeps: the key, what it needs to run, and its
|
||||
/// telemetry settings.
|
||||
const ALLOWED: &[&str] = &[
|
||||
const CLAUDE_ALLOWED: &[&str] = &[
|
||||
"ANTHROPIC_API_KEY",
|
||||
"PATH",
|
||||
"HOME",
|
||||
|
|
@ -29,37 +33,59 @@ const ALLOWED: &[&str] = &[
|
|||
"OTEL_RESOURCE_ATTRIBUTES",
|
||||
];
|
||||
|
||||
/// The variables of `vars` on [`ALLOWED`].
|
||||
pub fn allowed(vars: impl IntoIterator<Item = (OsString, OsString)>) -> Vec<(OsString, OsString)> {
|
||||
/// Every variable `opencode` keeps: the provider key, what it needs to run,
|
||||
/// and the config [`crate::acp`] renders for it.
|
||||
const OPENCODE_ALLOWED: &[&str] = &[
|
||||
crate::acp::PROVIDER_KEY_VAR,
|
||||
"PATH",
|
||||
"HOME",
|
||||
"OPENCODE_CONFIG",
|
||||
"OPENCODE_PERMISSION",
|
||||
"OPENCODE_DISABLE_PROJECT_CONFIG",
|
||||
];
|
||||
|
||||
/// The variables of `vars` on `allowed`.
|
||||
fn allowed(
|
||||
allowed: &[&str],
|
||||
vars: impl IntoIterator<Item = (OsString, OsString)>,
|
||||
) -> Vec<(OsString, OsString)> {
|
||||
vars.into_iter()
|
||||
.filter(|(name, _)| ALLOWED.iter().any(|a| name == a))
|
||||
.filter(|(name, _)| allowed.iter().any(|a| name == a))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// When this process is the trampoline, exec `claude` and never return.
|
||||
/// When this process is the trampoline, exec the real binary and never
|
||||
/// return.
|
||||
pub fn exec_if_requested() {
|
||||
let Some(claude) = std::env::var_os(EXEC_CLAUDE) else {
|
||||
return;
|
||||
};
|
||||
let error = Command::new(&claude)
|
||||
.args(std::env::args_os().skip(1))
|
||||
.env_clear()
|
||||
.envs(allowed(std::env::vars_os()))
|
||||
.exec();
|
||||
eprintln!("exec {}: {error}", claude.to_string_lossy());
|
||||
std::process::exit(127);
|
||||
for (var, allow) in [
|
||||
(EXEC_CLAUDE, CLAUDE_ALLOWED),
|
||||
(EXEC_OPENCODE, OPENCODE_ALLOWED),
|
||||
] {
|
||||
let Some(program) = std::env::var_os(var) else {
|
||||
continue;
|
||||
};
|
||||
let error = Command::new(&program)
|
||||
.args(std::env::args_os().skip(1))
|
||||
.env_clear()
|
||||
.envs(allowed(allow, std::env::vars_os()))
|
||||
.exec();
|
||||
eprintln!("exec {}: {error}", program.to_string_lossy());
|
||||
std::process::exit(127);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn only_the_allowlist_reaches_claude() {
|
||||
let daemon = [
|
||||
/// The daemon's environment on a run of either runtime.
|
||||
fn daemon() -> Vec<(OsString, OsString)> {
|
||||
[
|
||||
("ANTHROPIC_API_KEY", "k"),
|
||||
(crate::acp::PROVIDER_KEY_VAR, "p"),
|
||||
("HOME", "/tmp/run"),
|
||||
("OTEL_RESOURCE_ATTRIBUTES", "repo=o/r,pr=7"),
|
||||
("OPENCODE_PERMISSION", r#"{"*":"deny"}"#),
|
||||
(
|
||||
"BAO_CLIENT_KEY",
|
||||
"/run/credentials/x/review-classifier-bao-key",
|
||||
|
|
@ -71,15 +97,32 @@ mod tests {
|
|||
"tls://nats.example:4222",
|
||||
),
|
||||
(EXEC_CLAUDE, "/bin/claude"),
|
||||
(EXEC_OPENCODE, "/bin/opencode"),
|
||||
]
|
||||
.map(|(n, v)| (OsString::from(n), OsString::from(v)));
|
||||
let kept: Vec<String> = allowed(daemon)
|
||||
.map(|(n, v)| (OsString::from(n), OsString::from(v)))
|
||||
.into()
|
||||
}
|
||||
|
||||
fn kept(allow: &[&str]) -> Vec<String> {
|
||||
allowed(allow, daemon())
|
||||
.into_iter()
|
||||
.map(|(n, _)| n.to_string_lossy().into_owned())
|
||||
.collect();
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn only_the_allowlist_reaches_claude() {
|
||||
assert_eq!(
|
||||
kept,
|
||||
kept(CLAUDE_ALLOWED),
|
||||
["ANTHROPIC_API_KEY", "HOME", "OTEL_RESOURCE_ATTRIBUTES"]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn only_the_allowlist_reaches_opencode() {
|
||||
assert_eq!(
|
||||
kept(OPENCODE_ALLOWED),
|
||||
[crate::acp::PROVIDER_KEY_VAR, "HOME", "OPENCODE_PERMISSION"]
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
150
swarm-review-classifier/src/cost.rs
Normal file
150
swarm-review-classifier/src/cost.rs
Normal file
|
|
@ -0,0 +1,150 @@
|
|||
//! What a classification on the ACP runtime cost, exported as the counter
|
||||
//! `hyperhive.review_classifier.cost.usage` with the run's `repo`, `pr`,
|
||||
//! `model` and `currency`. On the claude runtime `claude` exports its own.
|
||||
//!
|
||||
//! The amount is computed here from the prices the run carries and the
|
||||
//! tokens the agent reports, so it does not depend on how, or whether, the
|
||||
//! agent prices a custom provider.
|
||||
//!
|
||||
//! The exporter is built on the first record, from the standard `OTEL_*`
|
||||
//! environment, as `hive-agent`'s `otel_turn_metrics` does; with no
|
||||
//! `OTEL_EXPORTER_OTLP_ENDPOINT` a record does nothing.
|
||||
|
||||
use std::sync::OnceLock;
|
||||
use std::time::Duration;
|
||||
|
||||
use hive_runtime::TokenUsage;
|
||||
use opentelemetry::KeyValue;
|
||||
use opentelemetry::metrics::{Counter, MeterProvider as _};
|
||||
use opentelemetry_otlp::{MetricExporter, Protocol, WithExportConfig};
|
||||
use opentelemetry_sdk::metrics::{PeriodicReader, SdkMeterProvider};
|
||||
use swarm_queue_client::review_classifier::Prices;
|
||||
|
||||
/// The currency every price is in.
|
||||
pub const CURRENCY: &str = "USD";
|
||||
|
||||
/// How often the counter is exported.
|
||||
const INTERVAL: Duration = Duration::from_mins(1);
|
||||
|
||||
/// One classification's cost, and what it is labelled with.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct Cost {
|
||||
pub repo: String,
|
||||
pub pr: u64,
|
||||
pub model: String,
|
||||
/// In [`CURRENCY`].
|
||||
pub amount: f64,
|
||||
}
|
||||
|
||||
/// Where a classification's cost goes: the metrics store in production, a
|
||||
/// list in the tests.
|
||||
pub trait CostSink: Send + Sync {
|
||||
fn record(&self, cost: &Cost);
|
||||
}
|
||||
|
||||
/// `usage` at `prices`. Cache reads and writes without a price of their own
|
||||
/// are priced as input.
|
||||
#[must_use]
|
||||
pub fn amount(prices: &Prices, usage: &TokenUsage) -> f64 {
|
||||
let per_token = |price: f64| price / 1_000_000.0;
|
||||
// One model call's count is far below `u32::MAX`.
|
||||
let tokens = |n: u64| f64::from(u32::try_from(n).unwrap_or(u32::MAX));
|
||||
tokens(usage.input_tokens) * per_token(prices.input)
|
||||
+ tokens(usage.output_tokens) * per_token(prices.output)
|
||||
+ tokens(usage.cache_read_input_tokens)
|
||||
* per_token(prices.cache_read.unwrap_or(prices.input))
|
||||
+ tokens(usage.cache_creation_input_tokens)
|
||||
* per_token(prices.cache_write.unwrap_or(prices.input))
|
||||
}
|
||||
|
||||
/// The counter in the metrics store.
|
||||
pub struct Otel;
|
||||
|
||||
impl CostSink for Otel {
|
||||
fn record(&self, cost: &Cost) {
|
||||
static COUNTER: OnceLock<Option<(SdkMeterProvider, Counter<f64>)>> = OnceLock::new();
|
||||
let Some((_, counter)) = COUNTER.get_or_init(build) else {
|
||||
return;
|
||||
};
|
||||
counter.add(
|
||||
cost.amount,
|
||||
&[
|
||||
KeyValue::new("repo", cost.repo.clone()),
|
||||
KeyValue::new("pr", i64::try_from(cost.pr).unwrap_or(i64::MAX)),
|
||||
KeyValue::new("model", cost.model.clone()),
|
||||
KeyValue::new("currency", CURRENCY),
|
||||
],
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// The provider, kept alive with the counter since it exports only while it
|
||||
/// lives, or `None` when no collector is configured or the exporter failed.
|
||||
fn build() -> Option<(SdkMeterProvider, Counter<f64>)> {
|
||||
if !std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT").is_ok_and(|s| !s.trim().is_empty()) {
|
||||
return None;
|
||||
}
|
||||
// The blocking client: `PeriodicReader` exports from a thread of its own
|
||||
// with no Tokio reactor. The resource comes from `OTEL_RESOURCE_ATTRIBUTES`.
|
||||
let exporter = MetricExporter::builder()
|
||||
.with_http()
|
||||
.with_protocol(Protocol::HttpJson)
|
||||
.build();
|
||||
let exporter = match exporter {
|
||||
Ok(exporter) => exporter,
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "building the cost metric exporter failed; no cost is exported");
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let reader = PeriodicReader::builder(exporter)
|
||||
.with_interval(INTERVAL)
|
||||
.build();
|
||||
let provider = SdkMeterProvider::builder().with_reader(reader).build();
|
||||
let counter = provider
|
||||
.meter("hyperhive.review_classifier")
|
||||
.f64_counter("hyperhive.review_classifier.cost.usage")
|
||||
.build();
|
||||
Some((provider, counter))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn usage(input: u64, output: u64, cache_read: u64, cache_write: u64) -> TokenUsage {
|
||||
let mut usage = TokenUsage::default();
|
||||
usage.input_tokens = input;
|
||||
usage.output_tokens = output;
|
||||
usage.cache_read_input_tokens = cache_read;
|
||||
usage.cache_creation_input_tokens = cache_write;
|
||||
usage
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn each_token_class_is_charged_at_its_price_per_million() {
|
||||
let prices = Prices {
|
||||
input: 3.0,
|
||||
output: 15.0,
|
||||
cache_read: Some(0.3),
|
||||
cache_write: Some(3.75),
|
||||
};
|
||||
let cost = amount(&prices, &usage(1_000, 200, 10_000, 2_000));
|
||||
assert!(
|
||||
(cost - (0.003 + 0.003 + 0.003 + 0.0075)).abs() < 1e-12,
|
||||
"{cost}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_without_a_price_of_its_own_is_charged_as_input() {
|
||||
let prices = Prices {
|
||||
input: 2.0,
|
||||
output: 0.0,
|
||||
cache_read: None,
|
||||
cache_write: None,
|
||||
};
|
||||
let cost = amount(&prices, &usage(0, 0, 500_000, 500_000));
|
||||
assert!((cost - 2.0).abs() < 1e-12, "{cost}");
|
||||
}
|
||||
}
|
||||
|
|
@ -7,7 +7,7 @@
|
|||
//! the controller reads the PR and sends it along.
|
||||
//!
|
||||
//! The store identity comes from the `BAO_*` environment. The queue client
|
||||
//! secret is read at start and again when authelia refuses it, the API key
|
||||
//! secret is read at start and again when authelia refuses it, the run's key
|
||||
//! once per run, both from the store under that identity.
|
||||
|
||||
use std::path::PathBuf;
|
||||
|
|
@ -19,9 +19,11 @@ use swarm_queue_client::review_classifier::{QUEUE_GROUP, Reply, SUBJECT, TryOut}
|
|||
use swarm_queue_client::stored_secret::StoredSecret;
|
||||
use swarm_queue_client::{ClientSecret, QueueConfig};
|
||||
|
||||
mod acp;
|
||||
mod aggregate;
|
||||
mod answer;
|
||||
mod child_env;
|
||||
mod cost;
|
||||
mod prompt;
|
||||
mod run;
|
||||
|
||||
|
|
@ -37,7 +39,8 @@ const PREFIX: &str = "SWARM_REVIEW_CLASSIFIER";
|
|||
const RETRY_DELAYS_S: [u64; 6] = [1, 2, 4, 8, 16, 30];
|
||||
|
||||
fn main() -> Result<()> {
|
||||
// Before anything else: as the trampoline, this process becomes `claude`.
|
||||
// Before anything else: as the trampoline, this process becomes the
|
||||
// model child.
|
||||
child_env::exec_if_requested();
|
||||
tokio::runtime::Builder::new_multi_thread()
|
||||
.enable_all()
|
||||
|
|
@ -54,10 +57,20 @@ async fn serve() -> Result<()> {
|
|||
std::env::var(&full).with_context(|| format!("{full} is unset"))
|
||||
};
|
||||
let client_id = var("OIDC_CLIENT_ID")?;
|
||||
let claude = run::Claude {
|
||||
program: PathBuf::from(var("CLAUDE")?),
|
||||
trampoline: std::env::current_exe()
|
||||
.context("finding this binary for the claude trampoline")?,
|
||||
let trampoline =
|
||||
std::env::current_exe().context("finding this binary for the child trampoline")?;
|
||||
let models = run::Runtimes {
|
||||
claude: run::Claude {
|
||||
program: PathBuf::from(var("CLAUDE")?),
|
||||
trampoline: trampoline.clone(),
|
||||
},
|
||||
acp: acp::Acp {
|
||||
program: PathBuf::from(var("OPENCODE")?),
|
||||
trampoline,
|
||||
deadline: swarm_queue_client::review_classifier::CLASSIFICATION_DEADLINE,
|
||||
cancel_wait: acp::CANCEL_WAIT,
|
||||
cost: cost::Otel,
|
||||
},
|
||||
};
|
||||
let secret = client_secret_with_retry(&client_id).await?;
|
||||
let reread_id = client_id.clone();
|
||||
|
|
@ -81,16 +94,16 @@ async fn serve() -> Result<()> {
|
|||
.with_context(|| format!("subscribing to {SUBJECT}"))?;
|
||||
tracing::info!(subject = SUBJECT, "review classifier: awaiting try-outs");
|
||||
|
||||
let claude = std::sync::Arc::new(claude);
|
||||
let models = std::sync::Arc::new(models);
|
||||
while let Some(msg) = requests.next().await {
|
||||
let Some(reply_to) = msg.reply.clone() else {
|
||||
tracing::warn!("a try-out arrived with no reply subject; dropping it");
|
||||
continue;
|
||||
};
|
||||
let queue = queue.clone();
|
||||
let claude = std::sync::Arc::clone(&claude);
|
||||
let models = std::sync::Arc::clone(&models);
|
||||
tokio::spawn(async move {
|
||||
let reply = handle(&msg.payload, &run::Store, &*claude).await;
|
||||
let reply = handle(&msg.payload, &run::Store, &*models).await;
|
||||
match serde_json::to_vec(&reply) {
|
||||
Ok(body) => {
|
||||
if let Err(e) = queue.publish(reply_to, body.into()).await {
|
||||
|
|
|
|||
|
|
@ -217,6 +217,7 @@ diff --git a/src/main.rs b/src/main.rs
|
|||
classifications: Vec::new(),
|
||||
model: "haiku".to_owned(),
|
||||
effort: None,
|
||||
runtime: swarm_queue_client::review_classifier::Runtime::ApiKey,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,11 +1,13 @@
|
|||
//! Answering one try-out: read the API key, ask the model one classification
|
||||
//! at a time, and turn the answers into the reply.
|
||||
//! Answering one try-out: read the key of the run's runtime, ask the model one
|
||||
//! classification at a time, and turn the answers into the reply.
|
||||
//!
|
||||
//! Each classification is one `claude --print` with no session, no tools and
|
||||
//! no MCP server, in a directory made for it and removed after it, so nothing
|
||||
//! carries from one run to the next. The API key is read from the store for
|
||||
//! every run and reaches only the child's environment, which keeps nothing of
|
||||
//! the daemon's beyond [`crate::child_env`]'s allowlist.
|
||||
//! On the API key runtime each classification is one `claude --print` with
|
||||
//! no session, no tools and no MCP server; on the ACP runtime it is one
|
||||
//! prompt on a new ACP session ([`crate::acp`]). Either way it runs in a
|
||||
//! directory made for it and removed after it, so nothing carries from one
|
||||
//! run to the next. The key is read from the store for every run and reaches
|
||||
//! only the child's environment, which keeps nothing of the daemon's beyond
|
||||
//! [`crate::child_env`]'s allowlist.
|
||||
|
||||
use std::future::Future;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
|
@ -15,33 +17,33 @@ use std::time::Duration;
|
|||
|
||||
use hive_claude::{Attach, Config, Sink};
|
||||
use swarm_queue_client::agent_inference::Effort;
|
||||
use swarm_queue_client::review_classifier::{CLASSIFICATION_DEADLINE, Reply, Row, TryOut};
|
||||
use swarm_secret_client::review_classifier::{ApiKey, api_key_path};
|
||||
use swarm_queue_client::review_classifier::{CLASSIFICATION_DEADLINE, Reply, Row, Runtime, TryOut};
|
||||
use swarm_secret_client::acp::ProviderKey;
|
||||
use swarm_secret_client::review_classifier::{ApiKey, acp_provider_key_path, api_key_path};
|
||||
|
||||
use crate::{aggregate, answer, prompt};
|
||||
|
||||
/// How long `claude` may print nothing before it is stopped.
|
||||
const IDLE_TIMEOUT: Duration = Duration::from_mins(2);
|
||||
|
||||
/// Where the run's API key comes from: the store in production, a fake in
|
||||
/// the tests.
|
||||
/// Where the run's key comes from: the store in production, a fake in the
|
||||
/// tests.
|
||||
pub trait KeySource {
|
||||
/// The key, or why there is none, in words the operator acts on.
|
||||
fn api_key(&self) -> impl Future<Output = Result<String, String>> + Send;
|
||||
/// The key `runtime` runs on, or why there is none, in words the
|
||||
/// operator acts on.
|
||||
fn key(&self, runtime: &Runtime) -> impl Future<Output = Result<String, String>> + Send;
|
||||
}
|
||||
|
||||
/// Asks the model one classification: `claude` in production, a fake in the
|
||||
/// tests.
|
||||
/// Asks the model one classification: [`Runtimes`] in production, a fake in
|
||||
/// the tests.
|
||||
pub trait Model {
|
||||
/// The final message of `model` at `effort` for `prompt`, or why there
|
||||
/// is none.
|
||||
/// The final message for `prompt` on `run`'s runtime, model and effort,
|
||||
/// run with `key`, or why there is none.
|
||||
fn ask(
|
||||
&self,
|
||||
prompt: &str,
|
||||
api_key: &str,
|
||||
resource_attributes: &str,
|
||||
model: &str,
|
||||
effort: Option<Effort>,
|
||||
key: &str,
|
||||
run: &TryOut,
|
||||
) -> impl Future<Output = Result<String, String>> + Send;
|
||||
}
|
||||
|
||||
|
|
@ -49,22 +51,15 @@ pub trait Model {
|
|||
/// run; otherwise each classification is a row, with an error where it has no
|
||||
/// answer.
|
||||
pub async fn answer(keys: &impl KeySource, model: &impl Model, run: &TryOut) -> Reply {
|
||||
let api_key = match keys.api_key().await {
|
||||
let key = match keys.key(&run.runtime).await {
|
||||
Ok(key) => key,
|
||||
Err(error) => return Reply::Failed { error },
|
||||
};
|
||||
let shared = prompt::shared_context(run);
|
||||
let attributes = resource_attributes(std::env::var("OTEL_RESOURCE_ATTRIBUTES").ok(), run);
|
||||
let mut rows: Vec<Row> = Vec::with_capacity(run.classifications.len());
|
||||
for c in &run.classifications {
|
||||
let outcome = model
|
||||
.ask(
|
||||
&prompt::classification_prompt(&shared, c),
|
||||
&api_key,
|
||||
&attributes,
|
||||
&run.model,
|
||||
run.effort,
|
||||
)
|
||||
.ask(&prompt::classification_prompt(&shared, c), &key, run)
|
||||
.await
|
||||
.and_then(|text| answer::parse(&text));
|
||||
rows.push(aggregate::row(c, outcome));
|
||||
|
|
@ -94,29 +89,72 @@ fn resource_attributes(base: Option<String>, run: &TryOut) -> String {
|
|||
}
|
||||
}
|
||||
|
||||
/// Each classification on the model its run's runtime names.
|
||||
pub struct Runtimes<C, A> {
|
||||
pub claude: C,
|
||||
pub acp: A,
|
||||
}
|
||||
|
||||
impl<C: Model + Sync, A: Model + Sync> Model for Runtimes<C, A> {
|
||||
async fn ask(&self, prompt: &str, key: &str, run: &TryOut) -> Result<String, String> {
|
||||
match run.runtime {
|
||||
Runtime::ApiKey => self.claude.ask(prompt, key, run).await,
|
||||
Runtime::Acp(_) => self.acp.ask(prompt, key, run).await,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The store, logged in to for each read: a run can come long after the last
|
||||
/// one's token expired.
|
||||
pub struct Store;
|
||||
|
||||
impl KeySource for Store {
|
||||
async fn api_key(&self) -> Result<String, String> {
|
||||
let path = api_key_path().map_err(|e| e.to_string())?;
|
||||
let store = swarm_secret_client::SecretStore::from_env(crate::CERT_ROLE)
|
||||
.await
|
||||
.map_err(|e| format!("logging in to the secret store to read the API key: {e}"))?;
|
||||
let stored: Option<ApiKey> = store
|
||||
.read_optional(&path)
|
||||
.await
|
||||
.map_err(|e| format!("reading secret/{path}: {e}"))?;
|
||||
match stored {
|
||||
Some(k) if !k.api_key.trim().is_empty() => Ok(k.api_key.trim().to_owned()),
|
||||
_ => Err(format!(
|
||||
"no API key at secret/{path}: write it in the secret store's UI, field `api_key`"
|
||||
)),
|
||||
async fn key(&self, runtime: &Runtime) -> Result<String, String> {
|
||||
match runtime {
|
||||
Runtime::ApiKey => {
|
||||
let path = api_key_path().map_err(|e| e.to_string())?;
|
||||
let key = read_key(&path, "API key", |k: ApiKey| k.api_key).await?;
|
||||
key.ok_or_else(|| {
|
||||
format!(
|
||||
"no API key at secret/{path}: write it in the secret store's UI, field \
|
||||
`api_key`"
|
||||
)
|
||||
})
|
||||
}
|
||||
Runtime::Acp(_) => {
|
||||
let path = acp_provider_key_path().map_err(|e| e.to_string())?;
|
||||
let key = read_key(&path, "ACP provider key", |k: ProviderKey| k.api_key).await?;
|
||||
key.ok_or_else(|| {
|
||||
format!(
|
||||
"no ACP provider key at secret/{path}: write it in the secret store's \
|
||||
UI, field `api_key`"
|
||||
)
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The key `field` takes from the object at `path`, trimmed; `None` when
|
||||
/// nothing or a blank is stored.
|
||||
async fn read_key<T: serde::de::DeserializeOwned>(
|
||||
path: &str,
|
||||
what: &str,
|
||||
field: impl FnOnce(T) -> String,
|
||||
) -> Result<Option<String>, String> {
|
||||
let store = swarm_secret_client::SecretStore::from_env(crate::CERT_ROLE)
|
||||
.await
|
||||
.map_err(|e| format!("logging in to the secret store to read the {what}: {e}"))?;
|
||||
let stored: Option<T> = store
|
||||
.read_optional(path)
|
||||
.await
|
||||
.map_err(|e| format!("reading secret/{path}: {e}"))?;
|
||||
Ok(stored
|
||||
.map(field)
|
||||
.map(|key| key.trim().to_owned())
|
||||
.filter(|key| !key.is_empty()))
|
||||
}
|
||||
|
||||
/// The `claude` binary the classifications run on, started through the
|
||||
/// trampoline that clears its environment.
|
||||
pub struct Claude {
|
||||
|
|
@ -127,17 +165,11 @@ pub struct Claude {
|
|||
}
|
||||
|
||||
impl Model for Claude {
|
||||
async fn ask(
|
||||
&self,
|
||||
prompt: &str,
|
||||
api_key: &str,
|
||||
attributes: &str,
|
||||
model: &str,
|
||||
effort: Option<Effort>,
|
||||
) -> Result<String, String> {
|
||||
async fn ask(&self, prompt: &str, api_key: &str, run: &TryOut) -> Result<String, String> {
|
||||
let attributes = resource_attributes(std::env::var("OTEL_RESOURCE_ATTRIBUTES").ok(), run);
|
||||
let dir = scratch_dir();
|
||||
let outcome = self
|
||||
.ask_in(&dir, prompt, api_key, attributes, model, effort)
|
||||
.ask_in(&dir, prompt, api_key, &attributes, &run.model, run.effort)
|
||||
.await;
|
||||
if let Err(e) = tokio::fs::remove_dir_all(&dir).await {
|
||||
tracing::warn!(dir = %dir.display(), error = %e, "removing a run's directory failed");
|
||||
|
|
@ -219,7 +251,7 @@ impl Claude {
|
|||
}
|
||||
|
||||
/// A directory of its own for one classification.
|
||||
fn scratch_dir() -> PathBuf {
|
||||
pub fn scratch_dir() -> PathBuf {
|
||||
static NEXT: AtomicU64 = AtomicU64::new(0);
|
||||
std::env::temp_dir().join(format!(
|
||||
"review-classifier-{}-{}",
|
||||
|
|
@ -268,11 +300,18 @@ mod tests {
|
|||
|
||||
use super::*;
|
||||
|
||||
struct Key(Result<&'static str, &'static str>);
|
||||
/// The API key runtime's key is `k` and the ACP runtime's `p`, or both
|
||||
/// fail with the same error.
|
||||
struct Key(Result<(), &'static str>);
|
||||
|
||||
impl KeySource for Key {
|
||||
async fn api_key(&self) -> Result<String, String> {
|
||||
self.0.map(str::to_owned).map_err(str::to_owned)
|
||||
async fn key(&self, runtime: &Runtime) -> Result<String, String> {
|
||||
self.0.map_err(str::to_owned)?;
|
||||
Ok(match runtime {
|
||||
Runtime::ApiKey => "k",
|
||||
Runtime::Acp(_) => "p",
|
||||
}
|
||||
.to_owned())
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -280,17 +319,16 @@ mod tests {
|
|||
struct Scripted;
|
||||
|
||||
impl Model for Scripted {
|
||||
async fn ask(
|
||||
&self,
|
||||
prompt: &str,
|
||||
api_key: &str,
|
||||
attrs: &str,
|
||||
model: &str,
|
||||
effort: Option<Effort>,
|
||||
) -> Result<String, String> {
|
||||
assert_eq!(api_key, "k");
|
||||
assert_eq!((model, effort), ("sonnet", Some(Effort::Low)));
|
||||
assert!(attrs.ends_with("repo=o/r,pr=7"), "{attrs}");
|
||||
async fn ask(&self, prompt: &str, key: &str, run: &TryOut) -> Result<String, String> {
|
||||
let wanted = match run.runtime {
|
||||
Runtime::ApiKey => "k",
|
||||
Runtime::Acp(_) => "p",
|
||||
};
|
||||
assert_eq!(key, wanted, "the key of the run's runtime");
|
||||
assert_eq!(
|
||||
(run.model.as_str(), run.effort),
|
||||
("sonnet", Some(Effort::Low))
|
||||
);
|
||||
if prompt.contains("ANSWER-FAIL") {
|
||||
Err("claude failed: exit 1".to_owned())
|
||||
} else if prompt.contains("ANSWER-MATCH") {
|
||||
|
|
@ -321,6 +359,23 @@ mod tests {
|
|||
.collect(),
|
||||
model: "sonnet".to_owned(),
|
||||
effort: Some(Effort::Low),
|
||||
runtime: Runtime::ApiKey,
|
||||
}
|
||||
}
|
||||
|
||||
fn acp_run(descriptions: &[&str]) -> TryOut {
|
||||
TryOut {
|
||||
runtime: Runtime::Acp(swarm_queue_client::review_classifier::Acp {
|
||||
base_url: "https://llm.example/v1".to_owned(),
|
||||
model: "qwen".to_owned(),
|
||||
prices: swarm_queue_client::review_classifier::Prices {
|
||||
input: 1.0,
|
||||
output: 1.0,
|
||||
cache_read: None,
|
||||
cache_write: None,
|
||||
},
|
||||
}),
|
||||
..run(descriptions)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -343,7 +398,7 @@ mod tests {
|
|||
#[tokio::test]
|
||||
async fn a_failed_or_malformed_answer_is_an_excluded_row_beside_the_others() {
|
||||
let reply = answer(
|
||||
&Key(Ok("k")),
|
||||
&Key(Ok(())),
|
||||
&Scripted,
|
||||
&run(&["ANSWER-MATCH", "ANSWER-FAIL", "ANSWER-JUNK"]),
|
||||
)
|
||||
|
|
@ -366,7 +421,7 @@ mod tests {
|
|||
#[tokio::test]
|
||||
async fn no_answer_at_all_fails_the_run_instead_of_approving() {
|
||||
let reply = answer(
|
||||
&Key(Ok("k")),
|
||||
&Key(Ok(())),
|
||||
&Scripted,
|
||||
&run(&["ANSWER-FAIL", "ANSWER-JUNK"]),
|
||||
)
|
||||
|
|
@ -377,6 +432,37 @@ mod tests {
|
|||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn an_acp_run_is_asked_with_the_acp_key() {
|
||||
let reply = answer(&Key(Ok(())), &Scripted, &acp_run(&["ANSWER-MATCH"])).await;
|
||||
assert!(matches!(reply, Reply::Done { .. }), "{reply:?}");
|
||||
}
|
||||
|
||||
/// Answers with its own name.
|
||||
struct Named(&'static str);
|
||||
|
||||
impl Model for Named {
|
||||
async fn ask(&self, _: &str, _: &str, _: &TryOut) -> Result<String, String> {
|
||||
Ok(self.0.to_owned())
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn each_run_goes_to_the_model_its_runtime_names() {
|
||||
let runtimes = Runtimes {
|
||||
claude: Named("claude"),
|
||||
acp: Named("acp"),
|
||||
};
|
||||
assert_eq!(
|
||||
runtimes.ask("p", "k", &run(&[])).await.as_deref(),
|
||||
Ok("claude")
|
||||
);
|
||||
assert_eq!(
|
||||
runtimes.ask("p", "k", &acp_run(&[])).await.as_deref(),
|
||||
Ok("acp")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn claude_is_started_through_the_trampoline() {
|
||||
let claude = Claude {
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
//! The trampoline through the real binary: a variable of the daemon's
|
||||
//! environment does not reach the program it execs, the API key does.
|
||||
//! environment does not reach the program it execs, the runtime's key does.
|
||||
|
||||
use std::process::Command;
|
||||
|
||||
|
|
@ -20,3 +20,26 @@ fn the_daemons_store_identity_does_not_reach_claude() {
|
|||
assert!(out.status.success(), "{out:?}");
|
||||
assert_eq!(String::from_utf8_lossy(&out.stdout), "unset|unset|k");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_daemons_store_identity_does_not_reach_opencode() {
|
||||
let out = Command::new(env!("CARGO_BIN_EXE_swarm-review-classifier"))
|
||||
// Stands in for `opencode`: prints what it was handed.
|
||||
.env("SWARM_REVIEW_CLASSIFIER_EXEC_OPENCODE", "/bin/sh")
|
||||
.env("BAO_CLIENT_KEY", "/run/credentials/x/review-classifier-bao-key")
|
||||
.env("CREDENTIALS_DIRECTORY", "/run/credentials/x")
|
||||
.env("ANTHROPIC_API_KEY", "k")
|
||||
.env("ACP_PROVIDER_API_KEY", "p")
|
||||
.env("OPENCODE_PERMISSION", r#"{"*":"deny"}"#)
|
||||
.args([
|
||||
"-c",
|
||||
r#"printf '%s|%s|%s|%s|%s' "${BAO_CLIENT_KEY-unset}" "${CREDENTIALS_DIRECTORY-unset}" "${ANTHROPIC_API_KEY-unset}" "${ACP_PROVIDER_API_KEY-unset}" "${OPENCODE_PERMISSION-unset}""#,
|
||||
])
|
||||
.output()
|
||||
.expect("the binary runs");
|
||||
assert!(out.status.success(), "{out:?}");
|
||||
assert_eq!(
|
||||
String::from_utf8_lossy(&out.stdout),
|
||||
r#"unset|unset|unset|p|{"*":"deny"}"#
|
||||
);
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue