Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 10 additions & 6 deletions .github/actions/setup-rust/action.yml
Original file line number Diff line number Diff line change
Expand Up @@ -125,12 +125,16 @@ runs:
# A copied rustup binary skips rustup-init, which normally creates env.
# Recreate that file after clearing our tool root: a runner login profile
# from an earlier install may still source this exact stable path.
cat > "$cargo_home/env" <<EOF
case ":\${PATH}:" in
*:"$cargo_home/bin":*) ;;
*) export PATH="$cargo_home/bin:\$PATH" ;;
esac
EOF
printf 'case ":${PATH}:" in\n *:"%s":*) ;;\n *) export PATH="%s:$PATH" ;;\nesac\n' \
"$cargo_home/bin" "$cargo_home/bin" > "$cargo_home/env"

# Earlier repository names left absolute Rust env sources in runner
# login profiles. Refresh only those owned entries to this active home.
if [[ "${GITHUB_ACTIONS:-}" == "true" ]]; then
bash .github/actions/setup-rust/repair-login-profile.sh \
"$cache_root" "$safe_toolchain" "$cargo_home" \
"$HOME/.profile" "$HOME/.bash_profile" "$HOME/.bash_login" "$HOME/.zprofile"
fi

echo "CARGO_HOME=$cargo_home" >> "$GITHUB_ENV"
echo "RUSTUP_HOME=$rustup_home" >> "$GITHUB_ENV"
Expand Down
41 changes: 41 additions & 0 deletions .github/actions/setup-rust/repair-login-profile.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
#!/usr/bin/env bash
set -euo pipefail

# Refresh only Rust setup's generated, quoted env sources for this toolchain
# and the two former Maestro repository names. Keep every other profile line.
cache_root="$1"
toolchain="$2"
cargo_home="$3"
active_env="$cargo_home/env"
shift 3

for profile_file in "$@"; do
[[ -e "$profile_file" ]] || continue
if [[ -L "$profile_file" || ! -f "$profile_file" ]]; then
echo '::error::Rust profile repair requires a regular, non-symlink profile.' >&2
exit 1
fi
scratch="$(mktemp "${profile_file}.maestro-rust.XXXXXX")"
trap 'rm -f "$scratch"' EXIT
cp -p "$profile_file" "$scratch"
: > "$scratch"
while IFS= read -r line || [[ -n "$line" ]]; do
for namespace in evalops-maestro dx-corp-maestro; do
old_env="$cache_root/$namespace/$toolchain/home/cargo/env"
# The dot/source command must be the whole line's command prefix. Do
# not rewrite comments, arbitrary strings, other repos, or toolchains.
case "$line" in
". \"$old_env\""*|"source \"$old_env\""*)
line="${line/\"$old_env\"/\"$active_env\"}"
;;
esac
done
printf '%s\n' "$line" >> "$scratch"
done < "$profile_file"
if ! cmp -s "$profile_file" "$scratch"; then
mv "$scratch" "$profile_file"
else
rm -f "$scratch"
fi
trap - EXIT
done
8 changes: 4 additions & 4 deletions .repository-projection.json
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,11 @@
"projection": "deixic-code",
"projectionSchemaVersion": 1,
"sourceRepository": "dx-corp/mono",
"sourceSha": "f818f42e9db2c906bda87b1a6f94f438898255b3",
"sourceSha": "083eacdbac0fef20630986791529393aa4e2683c",
"destinationRepository": "dx-corp/code",
"priorProjectedBase": "4880127235ca9993371266cdff567a35093a8985",
"definitionDigest": "82936441c776e3e8edb5d215a75007ec9714a233f489d460075d79d5ef5ba32f",
"priorProjectedBase": "402b1d7e517279b33d3f6d6aef9c77a0633a8266",
"definitionDigest": "cb9d429542ebb0a2de9b42a7aad60d9d8696a648ceba47c30f05c0b285ca0db7",
"toolDigest": "f8cb071b0f27267120ccf45a00d0982f45113bd23535bef6a1555b4933f99f13",
"contentDigest": "bd5022d7258195d203b9698c72f86a6b70182920f5c4aad17cce41c7ceee8b27",
"contentDigest": "b19cd9d7b8b22856b6565af0c70d61f3b0125ade7f5dd8e4494610dce7cd538d",
"publicationEligible": true
}
11 changes: 11 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

113 changes: 101 additions & 12 deletions packages/ai-rs/src/anthropic.rs
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,9 @@ struct SseParserState {
pending_stop_reason: Option<StopReason>,
/// Thinking signature from `signature_delta`, associated with thinking blocks
pending_thinking_signature: Option<String>,
usage_input_tokens: Option<u64>,
usage_cache_read_tokens: Option<u64>,
usage_cache_creation_tokens: Option<u64>,
}

const ANTHROPIC_DEFAULT_BASE_URL: &str = "https://api.anthropic.com/v1";
Expand Down Expand Up @@ -465,11 +468,21 @@ impl AnthropicClient {

// Add tools with cache_control on the last tool for tool definition caching
if !config.tools.is_empty() {
let mut tools_json: Vec<serde_json::Value> = config
.tools
.iter()
.map(|tool| {
let mut wire = serde_json::json!(tool);
if let Some(fields) = wire.as_object_mut() {
// Result declarations belong to discovery, not Anthropic's
// tool input contract (including its count-token endpoint).
fields.remove("output_schema");
}
wire
})
.collect();
if mark_tools {
// Add cache_control to the last tool for tool definition caching
let mut tools_json: Vec<serde_json::Value> =
config.tools.iter().map(|t| serde_json::json!(t)).collect();

// Add cache_control to the last tool
if let Some(last_tool) = tools_json.last_mut() {
if let Some(obj) = last_tool.as_object_mut() {
Expand All @@ -479,11 +492,8 @@ impl AnthropicClient {
);
}
}

body["tools"] = serde_json::json!(tools_json);
} else {
body["tools"] = serde_json::json!(config.tools.as_ref());
}
body["tools"] = serde_json::json!(tools_json);
}

if let Some(thinking) = &config.thinking {
Expand Down Expand Up @@ -634,6 +644,11 @@ fn parse_sse_event(data: &str, state: &mut SseParserState) -> Option<StreamEvent
match event_type {
"message_start" => {
let parsed: MessageStartEvent = serde_json::from_str(event_data).ok()?;
if let Some(usage) = &parsed.message.usage {
state.usage_input_tokens = usage.input_tokens;
state.usage_cache_read_tokens = usage.cache_read_input_tokens;
state.usage_cache_creation_tokens = usage.cache_creation_input_tokens;
}
Some(StreamEvent::MessageStart {
id: parsed.message.id,
model: parsed.message.model,
Expand Down Expand Up @@ -702,11 +717,19 @@ fn parse_sse_event(data: &str, state: &mut SseParserState) -> Option<StreamEvent
state.pending_stop_reason = delta.stop_reason;
}
}
parsed.usage.map(|usage| StreamEvent::Usage {
input_tokens: usage.input_tokens.unwrap_or(0),
output_tokens: usage.output_tokens.unwrap_or(0),
cache_read_tokens: usage.cache_read_input_tokens,
cache_creation_tokens: usage.cache_creation_input_tokens,
let usage = parsed.usage?;
state.usage_input_tokens = usage.input_tokens.or(state.usage_input_tokens);
state.usage_cache_read_tokens = usage
.cache_read_input_tokens
.or(state.usage_cache_read_tokens);
state.usage_cache_creation_tokens = usage
.cache_creation_input_tokens
.or(state.usage_cache_creation_tokens);
Some(StreamEvent::Usage {
input_tokens: state.usage_input_tokens?,
output_tokens: usage.output_tokens?,
cache_read_tokens: state.usage_cache_read_tokens,
cache_creation_tokens: state.usage_cache_creation_tokens,
})
}
"message_stop" => {
Expand Down Expand Up @@ -746,6 +769,8 @@ struct MessageStartEvent {
struct MessageInfo {
id: String,
model: String,
#[serde(default)]
usage: Option<UsageInfo>,
}

/// Content block start event - signals start of a new content block
Expand Down Expand Up @@ -885,6 +910,70 @@ mod tests {
SseParserState::default()
}

#[test]
fn partial_anthropic_usage_never_fabricates_completion_tokens() {
let mut state = new_state();
let partial = "event: message_delta\ndata: {\"type\":\"message_delta\",\"usage\":{\"input_tokens\":17}}";
assert!(!matches!(
parse_sse_event(partial, &mut state),
Some(StreamEvent::Usage { .. })
));
}

#[test]
fn discovery_output_schema_stays_out_of_anthropic_provider_tool_parameters() {
let schema = serde_json::json!({"type":"object","properties":{"label":{"type":"string"}},"required":["label"]});
let tool = Tool::new("classify", "Classify text").with_output_schema(schema.clone());
let messages = vec![Message {
role: Role::User,
content: MessageContent::Text("Classify a bird".into()),
}];
let client = AnthropicClient::new("test-key").unwrap();
for cached in [false, true] {
let config = RequestConfig {
model: "anthropic/claude-opus-4-7".into(),
tools: std::sync::Arc::new(vec![tool.clone()]),
cache_system_prompt: cached,
..Default::default()
};
for body in [
client.build_request_body(&messages, &config).unwrap(),
client.build_count_tokens_body(&messages, &config).unwrap(),
] {
let wire = &body["tools"][0];
assert_eq!(wire["name"], "classify");
assert_eq!(wire["input_schema"], tool.input_schema);
assert!(
wire.get("output_schema").is_none(),
"discovery metadata must not become an Anthropic tool parameter: {wire}"
);
}
}
assert_eq!(
tool.output_schema,
Some(schema),
"the admitted catalog retains the owner's output declaration"
);
}

#[test]
fn final_anthropic_usage_preserves_start_input_and_final_output() {
let mut state = new_state();
parse_sse_event(
"event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"id\",\"model\":\"claude\",\"usage\":{\"input_tokens\":17,\"output_tokens\":0}}}",
&mut state,
);
let event = parse_sse_event("event: message_delta\ndata: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":9},\"delta\":{\"stop_reason\":\"end_turn\"}}",&mut state).unwrap();
assert!(matches!(
event,
StreamEvent::Usage {
input_tokens: 17,
output_tokens: 9,
..
}
));
}

#[test]
fn parse_message_start() {
let data = r#"event: message_start
Expand Down
5 changes: 4 additions & 1 deletion packages/ai-rs/src/bedrock.rs
Original file line number Diff line number Diff line change
Expand Up @@ -460,7 +460,10 @@ fn build_content_block(block: &ContentBlock) -> Option<Result<bedrock::types::Co
}

fn build_image_block(source: &ImageSource) -> Result<bedrock::types::ContentBlock> {
let ImageSource::Base64 { media_type, data } = source else {
let ImageSource::Base64 {
media_type, data, ..
} = source
else {
bail!("Bedrock Converse requires base64 image sources");
};
let format = media_type
Expand Down
Loading
Loading