review classifier: classifications of one PR share a cached prefix
On the API key runtime the classifier now hands claude its prompt as one
stream-json user message (`--input-format stream-json`) with two text
blocks: the PR's shared context, carrying `cache_control: ephemeral`, and
the classification's question.
As plain stdin text the whole prompt was one block, and claude puts its
only message breakpoint on the last block, after the question. The cache
writes an entry only where a breakpoint is, so no entry ended at the
shared-context/question boundary. Every classification wrote the whole
prompt and none read it (production: cacheRead ~0, cacheCreation ~input).
The system-prompt prefix (~270 tokens) is under the minimum cacheable
length. Captures against a local listener with a dummy key (claude
2.1.223) show the client breakpoint survives on the shared block, and that
claude adds none there on its own.
The PR stays in the user role: moving it into the system prompt would give
untrusted PR text the system prompt's authority.
`Model::ask` takes a `Prompt { shared, question }`. The ACP runtime joins
it back into the same text it sent before.
Refs #5245
This commit is contained in:
parent
95b89970ee
commit
6ed88ad70a
5 changed files with 223 additions and 34 deletions
|
|
@ -14,7 +14,8 @@ forge and holds no forge credential.
|
|||
|
||||
- The shared context first: a system prompt, the PR's commits with their full
|
||||
messages, and the diff with lockfiles collapsed as `hive-forge diff` does.
|
||||
Every classification's prompt starts with the same bytes.
|
||||
Every classification's prompt starts with the same bytes, and `claude` gets
|
||||
that part as a content block of its own, with a cache breakpoint at its end.
|
||||
- One `claude --print` per classification, with no session, no tools
|
||||
(`--tools ""`), no MCP server, its own empty directory and
|
||||
`ANTHROPIC_API_KEY` in its environment only.
|
||||
|
|
|
|||
|
|
@ -78,12 +78,17 @@ pub struct Acp<S> {
|
|||
}
|
||||
|
||||
impl<S: CostSink> crate::run::Model for Acp<S> {
|
||||
async fn ask(&self, prompt: &str, key: &str, run: &TryOut) -> Result<String, String> {
|
||||
async fn ask(
|
||||
&self,
|
||||
prompt: &crate::prompt::Prompt<'_>,
|
||||
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;
|
||||
let outcome = self.ask_in(&dir, &prompt.text(), 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");
|
||||
}
|
||||
|
|
@ -275,6 +280,7 @@ mod tests {
|
|||
use swarm_queue_client::review_classifier::Prices;
|
||||
|
||||
use super::*;
|
||||
use crate::prompt::Prompt;
|
||||
use crate::run::Model as _;
|
||||
|
||||
/// An ACP agent in plain `sh`, logging to `<log>.*`. It writes its
|
||||
|
|
@ -388,7 +394,7 @@ done
|
|||
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;
|
||||
let answer = acp.ask(&Prompt::of("classify"), "p", &run()).await;
|
||||
assert_eq!(
|
||||
answer.as_deref(),
|
||||
Ok(r#"{"verdict":"match","certainty":90,"reason":"ok"}"#)
|
||||
|
|
@ -412,7 +418,9 @@ done
|
|||
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");
|
||||
acp.ask(&Prompt::of("classify"), "p", &run())
|
||||
.await
|
||||
.expect("an answer");
|
||||
let env = read(dir.path(), "env");
|
||||
for line in [
|
||||
"ACP_PROVIDER_API_KEY=p",
|
||||
|
|
@ -442,7 +450,9 @@ done
|
|||
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");
|
||||
acp.ask(&Prompt::of("classify"), "p", &run())
|
||||
.await
|
||||
.expect("an answer");
|
||||
let config: Value = serde_json::from_str(&read(dir.path(), "config")).expect("JSON");
|
||||
assert_eq!(config["default_agent"], AGENT_ID);
|
||||
// Only a prompt: a `permission` or `tools` key here would be merged
|
||||
|
|
@ -462,7 +472,9 @@ done
|
|||
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");
|
||||
acp.ask(&Prompt::of("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");
|
||||
|
|
@ -473,7 +485,7 @@ done
|
|||
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;
|
||||
let answer = acp.ask(&Prompt::of("STALL"), "p", &run()).await;
|
||||
assert_eq!(answer, Err("no answer within 0 seconds".to_owned()));
|
||||
assert!(
|
||||
read(dir.path(), "methods")
|
||||
|
|
@ -495,9 +507,12 @@ done
|
|||
);
|
||||
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");
|
||||
let answer = tokio::time::timeout(
|
||||
Duration::from_secs(10),
|
||||
acp.ask(&Prompt::of("one"), "p", &run()),
|
||||
)
|
||||
.await
|
||||
.expect("the deadline ends the ask");
|
||||
assert_eq!(answer, Err("no answer within 0 seconds".to_owned()));
|
||||
}
|
||||
|
||||
|
|
@ -506,7 +521,9 @@ done
|
|||
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");
|
||||
acp.ask(&Prompt::of("MUTE"), "p", &run())
|
||||
.await
|
||||
.expect("an answer");
|
||||
assert!(costs.0.lock().expect("unpoisoned").is_empty());
|
||||
}
|
||||
|
||||
|
|
@ -517,7 +534,11 @@ done
|
|||
let mut claude = run();
|
||||
claude.runtime = Runtime::ApiKey;
|
||||
|
||||
assert!(acp.ask("classify", "p", &claude).await.is_err());
|
||||
assert!(
|
||||
acp.ask(&Prompt::of("classify"), "p", &claude)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
assert_eq!(read(dir.path(), "methods"), "", "no agent started");
|
||||
}
|
||||
|
||||
|
|
@ -538,10 +559,16 @@ done
|
|||
let (acp, _) = acp(dir.path(), Duration::from_mins(1));
|
||||
|
||||
let (run_one, run_two) = (run(), run());
|
||||
let (one, two) = tokio::join!(acp.ask("one", "p", &run_one), acp.ask("two", "p", &run_two));
|
||||
let (prompt_one, prompt_two) = (Prompt::of("one"), Prompt::of("two"));
|
||||
let (one, two) = tokio::join!(
|
||||
acp.ask(&prompt_one, "p", &run_one),
|
||||
acp.ask(&prompt_two, "p", &run_two)
|
||||
);
|
||||
one.expect("an answer");
|
||||
two.expect("an answer");
|
||||
acp.ask("three", "p", &run()).await.expect("an answer");
|
||||
acp.ask(&Prompt::of("three"), "p", &run())
|
||||
.await
|
||||
.expect("an answer");
|
||||
let cwd = dir.path().join("cwd").display().to_string();
|
||||
assert_eq!(cwds(dir.path()), [cwd.clone(), cwd.clone(), cwd]);
|
||||
}
|
||||
|
|
@ -551,8 +578,12 @@ done
|
|||
let dir = tempfile::tempdir().expect("a dir");
|
||||
let (acp, _) = acp(dir.path(), Duration::from_mins(1));
|
||||
|
||||
acp.ask("one", "p", &run()).await.expect("an answer");
|
||||
acp.ask("two", "p", &run()).await.expect("an answer");
|
||||
acp.ask(&Prompt::of("one"), "p", &run())
|
||||
.await
|
||||
.expect("an answer");
|
||||
acp.ask(&Prompt::of("two"), "p", &run())
|
||||
.await
|
||||
.expect("an answer");
|
||||
let homes = read(dir.path(), "homes");
|
||||
let homes: Vec<&str> = homes.lines().collect();
|
||||
assert_eq!(homes.len(), 2, "{homes:?}");
|
||||
|
|
|
|||
|
|
@ -54,19 +54,46 @@ pub fn shared_context(run: &TryOut) -> String {
|
|||
out
|
||||
}
|
||||
|
||||
/// The prompt for one classification: `shared` unchanged, then the question.
|
||||
/// One classification's prompt: the run's shared context unchanged, then the
|
||||
/// question.
|
||||
pub struct Prompt<'a> {
|
||||
pub shared: &'a str,
|
||||
pub question: String,
|
||||
}
|
||||
|
||||
impl Prompt<'_> {
|
||||
/// The prompt as one text, for a runtime that takes no content blocks.
|
||||
#[must_use]
|
||||
pub fn text(&self) -> String {
|
||||
format!("{}{}", self.shared, self.question)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
impl Prompt<'static> {
|
||||
/// A prompt whose whole text is `text`.
|
||||
pub fn of(text: &str) -> Self {
|
||||
Self {
|
||||
shared: "",
|
||||
question: text.to_owned(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The prompt for one classification: `shared`, then the question.
|
||||
#[must_use]
|
||||
pub fn classification_prompt(shared: &str, c: &Classification) -> String {
|
||||
pub fn classification_prompt<'a>(shared: &'a str, c: &Classification) -> Prompt<'a> {
|
||||
let wanted = match c.polarity {
|
||||
Polarity::Positive => "a match is the wanted outcome",
|
||||
Polarity::Negative => "a match is a violation",
|
||||
};
|
||||
format!(
|
||||
"{shared}## The question\n\nClassification `{}` ({wanted}):\n\n{}\n\n\
|
||||
let question = format!(
|
||||
"## The question\n\nClassification `{}` ({wanted}):\n\n{}\n\n\
|
||||
Answer with the JSON object only. Keep the reason under {REASON_MAX_CHARS} characters.\n",
|
||||
c.name,
|
||||
c.description.trim()
|
||||
)
|
||||
);
|
||||
Prompt { shared, question }
|
||||
}
|
||||
|
||||
/// The lockfiles whose diffs collapse to one line. The same list as
|
||||
|
|
@ -259,10 +286,11 @@ diff --git a/src/main.rs b/src/main.rs
|
|||
]
|
||||
.map(|c| classification_prompt(&shared, &c));
|
||||
for prompt in &prompts {
|
||||
assert!(prompt.starts_with(&shared), "{prompt}");
|
||||
assert_eq!(prompt.shared, shared);
|
||||
assert!(prompt.text().starts_with(&shared), "{}", prompt.text());
|
||||
}
|
||||
// The question is all that differs, and none of it is in the prefix.
|
||||
assert_ne!(prompts[0], prompts[1]);
|
||||
assert_ne!(prompts[0].question, prompts[1].question);
|
||||
assert!(!shared.contains("datetime") && !shared.contains("no-ads"));
|
||||
// And the prefix is a function of the run alone.
|
||||
assert_eq!(shared, shared_context(&run()));
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ use swarm_queue_client::review_classifier::{CLASSIFICATION_DEADLINE, Reply, Row,
|
|||
use swarm_secret_client::acp::ProviderKey;
|
||||
use swarm_secret_client::review_classifier::{ApiKey, acp_provider_key_path, api_key_path};
|
||||
|
||||
use crate::prompt::Prompt;
|
||||
use crate::{aggregate, answer, prompt};
|
||||
|
||||
/// How long `claude` may print nothing before it is stopped.
|
||||
|
|
@ -41,7 +42,7 @@ pub trait Model {
|
|||
/// run with `key`, or why there is none.
|
||||
fn ask(
|
||||
&self,
|
||||
prompt: &str,
|
||||
prompt: &Prompt<'_>,
|
||||
key: &str,
|
||||
run: &TryOut,
|
||||
) -> impl Future<Output = Result<String, String>> + Send;
|
||||
|
|
@ -96,7 +97,7 @@ pub struct Runtimes<C, 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> {
|
||||
async fn ask(&self, prompt: &Prompt<'_>, 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,
|
||||
|
|
@ -165,7 +166,12 @@ pub struct Claude {
|
|||
}
|
||||
|
||||
impl Model for Claude {
|
||||
async fn ask(&self, prompt: &str, api_key: &str, run: &TryOut) -> Result<String, String> {
|
||||
async fn ask(
|
||||
&self,
|
||||
prompt: &Prompt<'_>,
|
||||
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
|
||||
|
|
@ -182,7 +188,7 @@ impl Claude {
|
|||
async fn ask_in(
|
||||
&self,
|
||||
dir: &Path,
|
||||
prompt: &str,
|
||||
prompt: &Prompt<'_>,
|
||||
api_key: &str,
|
||||
attributes: &str,
|
||||
model: &str,
|
||||
|
|
@ -198,7 +204,8 @@ impl Claude {
|
|||
hive_claude::Claude::spawn(&config, &Attach::OneOff).map_err(|e| e.to_string())?;
|
||||
let cancel = running.cancel_handle();
|
||||
let sink = ResultSink::default();
|
||||
let wait = running.wait(prompt, &sink);
|
||||
let message = stream_json_message(prompt);
|
||||
let wait = running.wait(&message, &sink);
|
||||
tokio::pin!(wait);
|
||||
match tokio::time::timeout(CLASSIFICATION_DEADLINE, &mut wait).await {
|
||||
Ok(Ok(())) => sink.take(),
|
||||
|
|
@ -233,6 +240,8 @@ impl Claude {
|
|||
strict_mcp_config: true,
|
||||
// `--tools ""`: no built-in tool at all.
|
||||
tools: Some(String::new()),
|
||||
// The prompt is [`stream_json_message`], not text.
|
||||
extra_args: vec!["--input-format".to_owned(), "stream-json".to_owned()],
|
||||
program: Some(self.trampoline.clone()),
|
||||
idle_timeout: Some(IDLE_TIMEOUT),
|
||||
claude_config_dir: Some(dir.join("config")),
|
||||
|
|
@ -250,6 +259,29 @@ impl Claude {
|
|||
}
|
||||
}
|
||||
|
||||
/// `prompt` as the one user message `--input-format stream-json` reads: the
|
||||
/// shared context and the question as separate text blocks, with a cache
|
||||
/// breakpoint on the shared one.
|
||||
///
|
||||
/// `claude` sets a breakpoint only on the last block, after the question, so
|
||||
/// without this one no cache entry ends where the run's classifications stop
|
||||
/// agreeing, and none of them reads the shared context from the cache. With
|
||||
/// the three `claude` sets (two in the system prompt, one on the question)
|
||||
/// this is four, the most one request may carry.
|
||||
fn stream_json_message(prompt: &Prompt<'_>) -> String {
|
||||
let message = serde_json::json!({
|
||||
"type": "user",
|
||||
"message": {
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": prompt.shared, "cache_control": {"type": "ephemeral"}},
|
||||
{"type": "text", "text": prompt.question},
|
||||
],
|
||||
},
|
||||
});
|
||||
format!("{message}\n")
|
||||
}
|
||||
|
||||
/// A directory of its own for one classification.
|
||||
pub fn scratch_dir() -> PathBuf {
|
||||
static NEXT: AtomicU64 = AtomicU64::new(0);
|
||||
|
|
@ -319,7 +351,13 @@ mod tests {
|
|||
struct Scripted;
|
||||
|
||||
impl Model for Scripted {
|
||||
async fn ask(&self, prompt: &str, key: &str, run: &TryOut) -> Result<String, String> {
|
||||
async fn ask(
|
||||
&self,
|
||||
prompt: &Prompt<'_>,
|
||||
key: &str,
|
||||
run: &TryOut,
|
||||
) -> Result<String, String> {
|
||||
let prompt = prompt.text();
|
||||
let wanted = match run.runtime {
|
||||
Runtime::ApiKey => "k",
|
||||
Runtime::Acp(_) => "p",
|
||||
|
|
@ -442,7 +480,7 @@ mod tests {
|
|||
struct Named(&'static str);
|
||||
|
||||
impl Model for Named {
|
||||
async fn ask(&self, _: &str, _: &str, _: &TryOut) -> Result<String, String> {
|
||||
async fn ask(&self, _: &Prompt<'_>, _: &str, _: &TryOut) -> Result<String, String> {
|
||||
Ok(self.0.to_owned())
|
||||
}
|
||||
}
|
||||
|
|
@ -454,11 +492,17 @@ mod tests {
|
|||
acp: Named("acp"),
|
||||
};
|
||||
assert_eq!(
|
||||
runtimes.ask("p", "k", &run(&[])).await.as_deref(),
|
||||
runtimes
|
||||
.ask(&Prompt::of("p"), "k", &run(&[]))
|
||||
.await
|
||||
.as_deref(),
|
||||
Ok("claude")
|
||||
);
|
||||
assert_eq!(
|
||||
runtimes.ask("p", "k", &acp_run(&[])).await.as_deref(),
|
||||
runtimes
|
||||
.ask(&Prompt::of("p"), "k", &acp_run(&[]))
|
||||
.await
|
||||
.as_deref(),
|
||||
Ok("acp")
|
||||
);
|
||||
}
|
||||
|
|
@ -524,4 +568,78 @@ mod tests {
|
|||
);
|
||||
assert_eq!(resource_attributes(None, &r), "repo=o/r,pr=7");
|
||||
}
|
||||
|
||||
/// The two content blocks of `line`, a stream-json user message.
|
||||
fn blocks(line: &str) -> Vec<serde_json::Value> {
|
||||
let message: serde_json::Value = serde_json::from_str(line).expect("one JSON message");
|
||||
assert_eq!(
|
||||
(&message["type"], &message["message"]["role"]),
|
||||
(&"user".into(), &"user".into())
|
||||
);
|
||||
let content = message["message"]["content"]
|
||||
.as_array()
|
||||
.expect("content blocks");
|
||||
assert_eq!(content.len(), 2, "{message}");
|
||||
content.clone()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_shared_context_is_its_own_block_with_the_cache_breakpoint() {
|
||||
let prompt = Prompt {
|
||||
shared: "# Pull request o/r#7\n\n## Diff\n",
|
||||
question: "## The question\n\nIs it?\n".to_owned(),
|
||||
};
|
||||
let line = stream_json_message(&prompt);
|
||||
assert!(
|
||||
line.ends_with('\n') && line.trim_end().lines().count() == 1,
|
||||
"{line}"
|
||||
);
|
||||
let [shared, question] = blocks(&line).try_into().expect("two blocks");
|
||||
assert_eq!(shared["text"], prompt.shared);
|
||||
assert_eq!(
|
||||
shared["cache_control"],
|
||||
serde_json::json!({"type": "ephemeral"})
|
||||
);
|
||||
assert_eq!(question["text"], prompt.question.as_str());
|
||||
assert!(question.get("cache_control").is_none(), "{question}");
|
||||
}
|
||||
|
||||
/// Stands in for `claude`: keeps its arguments and stdin beside itself,
|
||||
/// and answers.
|
||||
const CLAUDE: &str = r#"#!/bin/sh
|
||||
printf '%s\n' "$@" > "$0.args"
|
||||
cat > "$0.stdin"
|
||||
printf '{"type":"result","subtype":"success","is_error":false,"result":"ok"}\n'
|
||||
"#;
|
||||
|
||||
#[tokio::test]
|
||||
async fn claude_reads_the_shared_context_and_the_question_as_separate_blocks() {
|
||||
use std::os::unix::fs::PermissionsExt as _;
|
||||
|
||||
let tmp = tempfile::tempdir().expect("a dir");
|
||||
let script = tmp.path().join("claude.sh");
|
||||
std::fs::write(&script, CLAUDE).expect("written");
|
||||
std::fs::set_permissions(&script, std::fs::Permissions::from_mode(0o755)).expect("chmod");
|
||||
let claude = Claude {
|
||||
program: PathBuf::from("claude"),
|
||||
trampoline: script.clone(),
|
||||
};
|
||||
let prompt = Prompt {
|
||||
shared: "# Pull request o/r#7\n",
|
||||
question: "## The question\n".to_owned(),
|
||||
};
|
||||
let answer = claude
|
||||
.ask_in(&tmp.path().join("run"), &prompt, "k", "a=b", "sonnet", None)
|
||||
.await;
|
||||
assert_eq!(answer.as_deref(), Ok("ok"));
|
||||
let args = std::fs::read_to_string(tmp.path().join("claude.sh.args")).expect("args");
|
||||
assert!(args.contains("--input-format\nstream-json\n"), "{args}");
|
||||
let stdin = std::fs::read_to_string(tmp.path().join("claude.sh.stdin")).expect("stdin");
|
||||
let [shared, question] = blocks(&stdin).try_into().expect("two blocks");
|
||||
assert_eq!(
|
||||
(&shared["text"], &question["text"]),
|
||||
(&prompt.shared.into(), &prompt.question.as_str().into())
|
||||
);
|
||||
assert_eq!(shared["cache_control"]["type"], "ephemeral");
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue