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
|
|
@ -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