refactor(agent): simplify agent/llm types, error, fallback, register, skill and token modules

Modified agent/llm/types/capabilities.rs +2 −12 View agent/llm/types/capabilities.rs

@@ -1,9 +1,5 @@

11use serde::{Deserialize, Serialize};

22

3/// What a provider+model pair supports.

4///

5/// Serializable so capability tables can be declared in configuration rather

6/// than hard-coded per provider.

73#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]

84pub struct ProviderCapabilities {

95 #[serde(default)]

@@ -25,18 +21,12 @@

2521}

2622

2723impl ProviderCapabilities {

28 /// Whether `input_tokens` leaves room for `max_output_tokens` of output.

29 ///

30 /// A `max_context_tokens` of 0 means "unknown", and always fits — an

31 /// unconfigured capability table must not silently reject every request.

32 pub fn fits_context(&self, input_tokens: usize, output_tokens: usize) -> bool {

24 pub fn fits_context(&self, input_tokens: usize, output_tokens: usize) -> bool {

3325 self.max_context_tokens == 0

3426 || input_tokens.saturating_add(output_tokens) <= self.max_context_tokens

3527 }

3628

37 /// Remaining input budget after reserving room for `output_tokens`.

38 /// `None` when the context window is unknown.

39 pub fn remaining_input_budget(&self, output_tokens: usize) -> Option<usize> {

29 pub fn remaining_input_budget(&self, output_tokens: usize) -> Option<usize> {

4030 (self.max_context_tokens != 0)

4131 .then(|| self.max_context_tokens.saturating_sub(output_tokens))

4232 }

Modified agent/llm/types/message.rs +4 −14 View agent/llm/types/message.rs

@@ -47,9 +47,7 @@

4747 pub name: Option<String>,

4848 #[serde(default, skip_serializing_if = "Vec::is_empty")]

4949 pub tool_calls: Vec<ToolCall>,

50 /// Required when `role` is [`MessageRole::Tool`]; see

51 /// [`ChatMessage::validate`].

52 #[serde(default, skip_serializing_if = "Option::is_none")]

50 #[serde(default, skip_serializing_if = "Option::is_none")]

5351 pub tool_call_id: Option<String>,

5452}

5553

@@ -76,9 +74,7 @@

7674 Self::text(MessageRole::Assistant, text)

7775 }

7876

79 /// A tool result. `tool_call_id` is mandatory for this role, so this

80 /// constructor makes it impossible to build one that omits it.

81 pub fn tool(tool_call_id: impl Into<String>, text: impl Into<String>) -> Self {

77 pub fn tool(tool_call_id: impl Into<String>, text: impl Into<String>) -> Self {

8278 Self {

8379 tool_call_id: Some(tool_call_id.into()),

8480 ..Self::text(MessageRole::Tool, text)

@@ -95,11 +91,7 @@

9591 self

9692 }

9793

98 /// Check role-dependent invariants the type system cannot express.

99 ///

100 /// Deserialization bypasses the constructors, so anything arriving from

101 /// outside the process should be validated before it reaches a provider.

102 pub fn validate(&self) -> Result<(), &'static str> {

94 pub fn validate(&self) -> Result<(), &'static str> {

10395 match self.role {

10496 MessageRole::Tool if self.tool_call_id.is_none() => {

10597 Err("tool message requires tool_call_id")

@@ -113,9 +105,7 @@

113105 }

114106 }

115107

116 /// Plain-text view of the content, concatenating text parts and skipping

117 /// images.

118 pub fn text_content(&self) -> String {

108 pub fn text_content(&self) -> String {

119109 match &self.content {

120110 MessageContent::Text(s) => s.clone(),

121111 MessageContent::Parts(parts) => parts

Modified agent/llm/types/request.rs +39 −13 View agent/llm/types/request.rs

@@ -10,14 +10,7 @@

1010 Json,

1111}

1212

13/// A provider-agnostic chat completion request.

14///

15/// Serializable so requests can be recorded and replayed — the primary way to

16/// build regression tests against real provider traffic.

17///

1813/// `#[non_exhaustive]` keeps future fields from being breaking changes; build

19/// values with [`ChatRequest::new`] plus the `with_*` methods, or from

20/// [`Default`].

2114#[derive(Debug, Clone, Default, Serialize, Deserialize)]

2215#[non_exhaustive]

2316pub struct ChatRequest {

@@ -58,16 +51,18 @@

5851 self

5952 }

6053

61 /// Set sampling temperature, clamped to the `0.0..=2.0` range every major

62 /// provider accepts. Out-of-range values are otherwise only rejected at

63 /// the API boundary, after the request has been paid for.

64 pub fn with_temperature(mut self, temperature: f32) -> Self {

54 pub fn with_temperature(mut self, temperature: f32) -> Self {

55 if !temperature.is_finite() {

56 return self;

57 }

6558 self.temperature = Some(temperature.clamp(0.0, 2.0));

6659 self

6760 }

6861

69 /// Set nucleus sampling cutoff, clamped to `0.0..=1.0`.

70 pub fn with_top_p(mut self, top_p: f32) -> Self {

62 pub fn with_top_p(mut self, top_p: f32) -> Self {

63 if !top_p.is_finite() {

64 return self;

65 }

7166 self.top_p = Some(top_p.clamp(0.0, 1.0));

7267 self

7368 }

@@ -120,6 +115,37 @@

120115 }

121116

122117 #[test]

118 fn nan_and_infinity_inputs_are_dropped() {

119 let req = ChatRequest::default()

120 .with_temperature(f32::NAN)

121 .with_top_p(f32::INFINITY);

122 assert_eq!(req.temperature, None);

123 assert_eq!(req.top_p, None);

124

125 let req = ChatRequest::default()

126 .with_temperature(f32::NEG_INFINITY)

127 .with_top_p(f32::NAN);

128 assert_eq!(req.temperature, None);

129 assert_eq!(req.top_p, None);

130 }

131

132 #[test]

133 fn finite_inputs_still_replace_previous_value() {

134 let req = ChatRequest::default()

135 .with_temperature(0.7)

136 .with_temperature(f32::NAN);

137 assert_eq!(req.temperature, Some(0.7));

138 }

139

140 #[test]

141 fn sampling_parameters_clamp_inclusive_endpoints() {

142 assert_eq!(ChatRequest::default().with_temperature(0.0).temperature, Some(0.0));

143 assert_eq!(ChatRequest::default().with_temperature(2.0).temperature, Some(2.0));

144 assert_eq!(ChatRequest::default().with_top_p(0.0).top_p, Some(0.0));

145 assert_eq!(ChatRequest::default().with_top_p(1.0).top_p, Some(1.0));

146 }

147

148 #[test]

123149 fn round_trips_through_json() {

124150 let req = ChatRequest::new(vec![ChatMessage::user("hi")]).with_max_tokens(8);

125151 let json = serde_json::to_string(&req).unwrap();

Modified agent/llm/types/response.rs +2 −9 View agent/llm/types/response.rs

@@ -2,10 +2,6 @@

22use crate::types::tools::ToolCall;

33use serde::{Deserialize, Serialize};

44

5/// A completed (non-streaming) chat response.

6///

7/// Serializable so responses can be recorded alongside their requests and

8/// replayed in tests.

95#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]

106pub struct ChatResponse {

117 pub content: String,

@@ -50,14 +46,11 @@

5046 self

5147 }

5248

53 /// Whether the model asked to invoke tools, meaning the agent loop should

54 /// execute them and continue rather than treat `content` as final.

55 pub fn requests_tools(&self) -> bool {

49 pub fn requests_tools(&self) -> bool {

5650 !self.tool_calls.is_empty()

5751 }

5852

59 /// Whether the response was cut short and `content` may be incomplete.

60 pub fn is_truncated(&self) -> bool {

53 pub fn is_truncated(&self) -> bool {

6154 matches!(

6255 self.finish_reason,

6356 FinishReason::MaxTokens | FinishReason::ContentFilter | FinishReason::Error

Modified agent/llm/types/stream.rs +10 −19 View agent/llm/types/stream.rs

@@ -24,15 +24,6 @@

2424 },

2525}

2626

27/// Reassembles a streamed response into its final form.

28///

29/// Providers emit tool calls as fragments — the arguments arrive as partial

30/// JSON text spread over many events, and `name` typically appears only on the

31/// first fragment. Every provider would otherwise reimplement this same

32/// grouping/concatenation/parse state machine, so it lives here once.

33///

34/// Fragments are grouped by `index` when the provider supplies one (the

35/// reliable discriminator for parallel tool calls) and by `id` otherwise.

3627#[derive(Debug, Default)]

3728pub struct StreamAccumulator {

3829 text: String,

@@ -75,8 +66,6 @@

7566 None => CallKey::Id(id.clone()),

7667 };

7768 let entry = self.calls.entry(key).or_default();

78 // Later fragments usually carry an empty id/name; never let

79 // them erase what the opening fragment established.

8069 if !id.is_empty() {

8170 entry.id = id.clone();

8271 }

@@ -111,11 +100,7 @@

111100 self.finish_reason

112101 }

113102

114 /// Parse the accumulated tool-call fragments.

115 ///

116 /// Fails if a call never received a name or if its arguments did not form

117 /// valid JSON — a truncated stream is a real failure, not an empty call.

118 pub fn tool_calls(&self) -> Result<Vec<ToolCall>, LlmError> {

103 pub fn tool_calls(&self) -> Result<Vec<ToolCall>, LlmError> {

119104 self.calls

120105 .values()

121106 .map(|p| {

@@ -139,9 +124,7 @@

139124 .collect()

140125 }

141126

142 /// Build the final response. `Done` must have been observed, otherwise the

143 /// stream ended early and the result would silently look complete.

144 pub fn finish(

127 pub fn finish(

145128 self,

146129 provider: impl Into<String>,

147130 model: impl Into<String>,

@@ -252,6 +235,14 @@

252235 }

253236

254237 #[test]

238 fn mixed_index_and_id_keys_for_same_call_is_malformed() {

239 let mut acc = StreamAccumulator::new();

240 acc.push(&delta(Some(0), "c1", Some("search"), r#"{"q":"ru"#));

241 acc.push(&delta(None, "c1", None, r#"st"}"#));

242 assert!(matches!(acc.tool_calls(), Err(LlmError::Stream(_))));

243 }

244

245 #[test]

255246 fn usage_events_accumulate() {

256247 let mut acc = StreamAccumulator::new();

257248 acc.push(&StreamEvent::Usage(TokenUsageStats::new(5, 1)));

Modified agent/llm/types/tools.rs +4 −9 View agent/llm/types/tools.rs

@@ -7,25 +7,20 @@

77 pub arguments: serde_json::Value,

88}

99

10/// How a tool may be scheduled relative to other tool calls in the same turn.

1110#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]

1211#[serde(rename_all = "snake_case")]

1312pub enum ConcurrencyPolicy {

14 /// Safe to run in parallel with any other `Safe` tool.

15 #[default]

13 #[default]

1614 Safe,

17 /// Must run alone, with no other tool in flight.

18 Exclusive,

19 /// Must run in the order the model emitted it, but may overlap `Safe` work.

20 Ordered,

15 Exclusive,

16 Ordered,

2117}

2218

2319#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]

2420pub struct ToolDefinition {

2521 pub name: String,

2622 pub description: String,

27 /// JSON Schema for the tool's arguments.

28 pub parameters: serde_json::Value,

23 pub parameters: serde_json::Value,

2924 #[serde(default)]

3025 pub concurrency: ConcurrencyPolicy,

3126}

Modified agent/llm/error.rs +53 −11 View agent/llm/error.rs

@@ -73,16 +73,16 @@

7373

7474 #[error("Skill rendering failed: {0}")]

7575 SkillRenderFailed(String),

76

77 #[error("primary failed (retryable), then fallback also failed: primary={primary} -> fallback={fallback}")]

78 FallbackExhausted {

79 primary: Box<LlmError>,

80 fallback: Box<LlmError>,

81 },

7682}

7783

7884impl LlmError {

79 /// Whether retrying the same request against another provider (or the same

80 /// one after a delay) could plausibly succeed.

81 ///

82 /// Deliberately excludes `ContextLengthExceeded`, `ContentFiltered`, and

83 /// `AuthError`: those reproduce deterministically, so retrying only burns

84 /// quota. `ApiError` is retryable for 5xx and 429 only.

85 pub fn is_retryable(&self) -> bool {

85 pub fn is_retryable(&self) -> bool {

8686 match self {

8787 Self::RateLimited { .. } | Self::Network(_) => true,

8888 Self::ApiError { status, .. } => *status == 429 || (500..600).contains(status),

@@ -94,22 +94,22 @@

9494 | std::io::ErrorKind::ConnectionAborted

9595 | std::io::ErrorKind::BrokenPipe

9696 ),

97 Self::FallbackExhausted { .. } => false,

9798 _ => false,

9899 }

99100 }

100101

101 /// Provider-advertised delay before retrying, when one was supplied.

102 pub fn retry_after(&self) -> Option<Duration> {

102 pub fn retry_after(&self) -> Option<Duration> {

103103 match self {

104104 Self::RateLimited {

105105 retry_after_secs, ..

106106 } => retry_after_secs.map(Duration::from_secs),

107 Self::FallbackExhausted { fallback, .. } => fallback.retry_after(),

107108 _ => None,

108109 }

109110 }

110111

111 /// The provider that produced this error, when attributable.

112 pub fn provider(&self) -> Option<&str> {

112 pub fn provider(&self) -> Option<&str> {

113113 match self {

114114 Self::RateLimited { provider, .. }

115115 | Self::ContentFiltered { provider, .. }

@@ -117,6 +117,7 @@

117117 | Self::ModelNotFound { provider, .. }

118118 | Self::ApiError { provider, .. }

119119 | Self::ProviderNotFound(provider) => Some(provider),

120 Self::FallbackExhausted { fallback, .. } => fallback.provider(),

120121 _ => None,

121122 }

122123 }

@@ -197,4 +198,45 @@

197198 );

198199 assert_eq!(LlmError::Stream("x".into()).provider(), None);

199200 }

201

202 #[test]

203 fn fallback_exhausted_is_not_retryable() {

204 let err = LlmError::FallbackExhausted {

205 primary: Box::new(LlmError::RateLimited {

206 provider: "primary".into(),

207 retry_after_secs: Some(1),

208 }),

209 fallback: Box::new(LlmError::Network("down".into())),

210 };

211 assert!(!err.is_retryable());

212 }

213

214 #[test]

215 fn fallback_exhausted_passes_through_fallback_metadata() {

216 let err = LlmError::FallbackExhausted {

217 primary: Box::new(LlmError::RateLimited {

218 provider: "primary".into(),

219 retry_after_secs: Some(1),

220 }),

221 fallback: Box::new(LlmError::ApiError {

222 provider: "fb".into(),

223 status: 503,

224 message: "down".into(),

225 }),

226 };

227 assert_eq!(err.retry_after(), None);

228 assert_eq!(err.provider(), Some("fb"));

229 }

230

231 #[test]

232 fn fallback_exhausted_propagates_fallback_retry_after() {

233 let err = LlmError::FallbackExhausted {

234 primary: Box::new(LlmError::Network("reset".into())),

235 fallback: Box::new(LlmError::RateLimited {

236 provider: "fb".into(),

237 retry_after_secs: Some(7),

238 }),

239 };

240 assert_eq!(err.retry_after(), Some(Duration::from_secs(7)));

241 }

200242}

Modified agent/llm/fallback.rs +93 −24 View agent/llm/fallback.rs

@@ -13,10 +13,6 @@

1313 pub model: String,

1414}

1515

16/// Routes to a secondary provider when the primary fails retryably.

17///

18/// Failover triggers on [`LlmError::is_retryable`] rather than on rate limits

19/// alone, so transient network faults and provider 5xx are covered too.

2016pub struct FallbackProvider {

2117 primary: Arc<dyn LlmProvider>,

2218 fallback: Option<FallbackTarget>,

@@ -27,12 +23,7 @@

2723 Self { primary, fallback }

2824 }

2925

30 /// Point a request at the fallback model.

31 ///

32 /// An explicit `model_override` from the caller is replaced — the fallback

33 /// provider almost certainly cannot serve the primary's model name — but

34 /// the discarded value is logged so the substitution is never silent.

35 fn retarget(mut request: ChatRequest, model: &str) -> ChatRequest {

26 fn retarget(mut request: ChatRequest, model: &str) -> ChatRequest {

3627 if let Some(previous) = &request.model_override

3728 && previous != model

3829 {

@@ -70,8 +61,6 @@

7061 }

7162

7263 fn capabilities(&self) -> ProviderCapabilities {

73 // Intersect: advertising a capability the fallback lacks would let a

74 // request be built that the fallback cannot serve.

7564 let Some(fb) = &self.fallback else {

7665 return self.primary.capabilities();

7766 };

@@ -95,7 +84,13 @@

9584 match self.primary.chat(request.clone()).await {

9685 Err(err) if err.is_retryable() => {

9786 self.log_failover(fb, &err, "request");

98 fb.provider.chat(Self::retarget(request, &fb.model)).await

87 fb.provider

88 .chat(Self::retarget(request, &fb.model))

89 .await

90 .map_err(|fallback_err| LlmError::FallbackExhausted {

91 primary: Box::new(err),

92 fallback: Box::new(fallback_err),

93 })

9994 }

10095 other => other,

10196 }

@@ -112,14 +107,15 @@

112107 return fb

113108 .provider

114109 .chat_stream(Self::retarget(request, &fb.model))

115 .await;

110 .await

111 .map_err(|fallback_err| LlmError::FallbackExhausted {

112 primary: Box::new(err),

113 fallback: Box::new(fallback_err),

114 });

116115 }

117116 other => other?,

118117 };

119118

120 // A stream can also fail *after* opening. Peek one event: until

121 // something has been yielded to the caller, a retryable error is still

122 // safely recoverable by restarting on the fallback.

123119 let Some(first) = stream.next().await else {

124120 return Ok(Box::pin(tokio_stream::empty()));

125121 };

@@ -130,8 +126,11 @@

130126 fb.provider

131127 .chat_stream(Self::retarget(request, &fb.model))

132128 .await

129 .map_err(|fallback_err| LlmError::FallbackExhausted {

130 primary: Box::new(err),

131 fallback: Box::new(fallback_err),

132 })

133133 }

134 // Put the peeked event back at the head of the stream.

135134 first => Ok(Box::pin(

136135 futures::stream::once(async move { first }).chain(stream),

137136 )),

@@ -150,15 +149,12 @@

150149 use std::sync::Mutex;

151150 use std::sync::atomic::{AtomicUsize, Ordering};

152151

153 /// A provider whose responses are scripted per call.

154 struct Scripted {

152 struct Scripted {

155153 id: String,

156154 caps: ProviderCapabilities,

157155 calls: AtomicUsize,

158 /// Error to return from `chat`/`chat_stream`, if any.

159 fail_with: Mutex<Option<LlmError>>,

160 /// When set, `chat_stream` opens successfully then yields this error.

161 fail_mid_stream: Mutex<Option<LlmError>>,

156 fail_with: Mutex<Option<LlmError>>,

157 fail_mid_stream: Mutex<Option<LlmError>>,

162158 seen_model_override: Mutex<Option<String>>,

163159 }

164160

@@ -384,4 +380,77 @@

384380 Err(LlmError::RateLimited { .. })

385381 ));

386382 }

383

384 #[tokio::test]

385 async fn fallback_failure_chains_primary_error() {

386 let secondary = Arc::new(Scripted::failing(

387 "secondary",

388 LlmError::ApiError {

389 provider: "secondary".into(),

390 status: 503,

391 message: "down".into(),

392 },

393 ));

394 let fb = FallbackProvider::new(

395 Arc::new(Scripted::failing("primary", rate_limited())),

396 Some(target(secondary.clone())),

397 );

398 let err = fb.chat(request()).await.expect_err("both must fail");

399 match err {

400 LlmError::FallbackExhausted { primary, fallback } => {

401 assert!(matches!(*primary, LlmError::RateLimited { .. }));

402 assert!(matches!(

403 *fallback,

404 LlmError::ApiError {

405 status: 503,

406 ..

407 }

408 ));

409 }

410 other => panic!("expected FallbackExhausted, got {other:?}"),

411 }

412 assert_eq!(secondary.calls.load(Ordering::SeqCst), 1);

413 }

414

415 #[tokio::test]

416 async fn stream_fallback_failure_chains_primary_error() {

417 let secondary = Arc::new(Scripted::failing(

418 "secondary",

419 LlmError::Network("unreachable".into()),

420 ));

421 let fb = FallbackProvider::new(

422 Arc::new(Scripted::failing_mid_stream("primary", rate_limited())),

423 Some(target(secondary.clone())),

424 );

425 let result = fb.chat_stream(request()).await;

426 let err = match result {

427 Ok(_) => panic!("expected FallbackExhausted, got Ok"),

428 Err(e) => e,

429 };

430 assert!(

431 matches!(err, LlmError::FallbackExhausted { .. }),

432 "expected FallbackExhausted, got {err:?}"

433 );

434 }

435

436 #[tokio::test]

437 async fn deterministic_failure_is_not_wrapped_in_fallback_exhausted() {

438 let secondary = Arc::new(Scripted::new("secondary"));

439 let fb = FallbackProvider::new(

440 Arc::new(Scripted::failing(

441 "primary",

442 LlmError::ContextLengthExceeded {

443 tokens_used: 10,

444 limit: 5,

445 },

446 )),

447 Some(target(secondary.clone())),

448 );

449 let surfaced = fb.chat(request()).await.expect_err("must fail");

450 assert!(

451 !matches!(surfaced, LlmError::FallbackExhausted { .. }),

452 "non-retryable primary error must not be wrapped, got {surfaced:?}"

453 );

454 assert_eq!(secondary.calls.load(Ordering::SeqCst), 0);

455 }

387456}

Modified agent/llm/provider.rs +2 −14 View agent/llm/provider.rs

@@ -22,16 +22,7 @@

2222 request: ChatRequest,

2323 ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamEvent, LlmError>> + Send>>, LlmError>;

2424

25 /// Estimate the input tokens this request will consume.

26 ///

27 /// The default implementation covers messages, tool-call payloads, and the

28 /// tool schemas — tool definitions are real input tokens and routinely run

29 /// to thousands, so omitting them badly under-counts tool-heavy requests.

30 ///

31 /// This is a script-aware heuristic, not a tokenizer: it is suitable for

32 /// budgeting, routing, and logging, but not for billing. Providers that

33 /// ship a real tokenizer should override it.

34 fn count_tokens(&self, request: &ChatRequest) -> usize {

25 fn count_tokens(&self, request: &ChatRequest) -> usize {

3526 let messages: usize = request

3627 .messages

3728 .iter()

@@ -42,9 +33,6 @@

4233 .iter()

4334 .map(|p| match p {

4435 ContentPart::Text(s) => estimate_text_tokens(s),

45 // Real cost scales with pixel dimensions, which a

46 // URL does not reveal; over-estimate rather than

47 // overflow the context at the provider.

4836 ContentPart::Image { detail, .. } => match detail {

4937 ImageDetail::Low => IMAGE_TOKENS_LOW,

5038 ImageDetail::High | ImageDetail::Auto => IMAGE_TOKENS_HIGH,

@@ -151,7 +139,7 @@

151139

152140 #[test]

153141 fn cjk_is_not_inflated_by_byte_counting() {

154 let zh = "这是一段中文测试文本内容"; // 12 chars / 36 bytes

142 let zh = "这是一段中文测试文本内容";

155143 let req = ChatRequest::new(vec![user(zh)]);

156144 assert_eq!(Stub.count_tokens(&req), 12 + MESSAGE_OVERHEAD_TOKENS);

157145 }

Modified agent/llm/register.rs +15 −22 View agent/llm/register.rs

@@ -2,10 +2,6 @@

22use std::collections::HashMap;

33use std::sync::Arc;

44

5/// Identifies a concrete provider+model pair.

6///

7/// Keying the registry on `provider_id` alone would let `openai/gpt-4o` and

8/// `openai/o3` silently displace one another, so both halves participate.

95#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]

106pub struct ProviderKey {

117 pub provider_id: String,

@@ -39,28 +35,27 @@

3935 }

4036 }

4137

42 /// Register a provider under its `(provider_id, model_id)` pair, returning

43 /// the provider it displaced, if any.

44 pub fn register(&mut self, provider: Arc<dyn LlmProvider>) -> Option<Arc<dyn LlmProvider>> {

38 /// Displacement is logged at warn level: most call sites ignore the

39 pub fn register(&mut self, provider: Arc<dyn LlmProvider>) -> Option<Arc<dyn LlmProvider>> {

4540 let key = ProviderKey::new(provider.provider_id(), provider.model_id());

46 self.providers.insert(key, provider)

41 let displaced = self.providers.insert(key.clone(), provider);

42 if displaced.is_some() {

43 tracing::warn!(

44 provider_id = %key.provider_id,

45 model_id = %key.model_id,

46 "provider re-registered; previous instance displaced",

47 );

48 }

49 displaced

4750 }

4851

49 /// Look up an exact provider+model pair.

50 ///

51 /// Returns an owned `Arc` rather than a borrow so callers can hold it

52 /// across an `.await` without keeping the registry borrowed.

53 pub fn get(&self, provider_id: &str, model_id: &str) -> Option<Arc<dyn LlmProvider>> {

52 pub fn get(&self, provider_id: &str, model_id: &str) -> Option<Arc<dyn LlmProvider>> {

5453 self.providers

5554 .get(&ProviderKey::new(provider_id, model_id))

5655 .cloned()

5756 }

5857

59 /// Any registered model for `provider_id`.

60 ///

61 /// Ties are broken by `model_id` ordering so repeated calls agree with each

62 /// other rather than following `HashMap` iteration order.

63 pub fn any_for_provider(&self, provider_id: &str) -> Option<Arc<dyn LlmProvider>> {

58 pub fn any_for_provider(&self, provider_id: &str) -> Option<Arc<dyn LlmProvider>> {

6459 self.providers

6560 .iter()

6661 .filter(|(k, _)| k.provider_id == provider_id)

@@ -78,15 +73,13 @@

7873 .remove(&ProviderKey::new(provider_id, model_id))

7974 }

8075

81 /// All registered `(provider_id, model_id)` pairs, sorted for stable output.

82 pub fn keys(&self) -> Vec<ProviderKey> {

76 pub fn keys(&self) -> Vec<ProviderKey> {

8377 let mut keys: Vec<_> = self.providers.keys().cloned().collect();

8478 keys.sort();

8579 keys

8680 }

8781

88 /// Distinct provider ids, sorted and deduplicated.

89 pub fn provider_ids(&self) -> Vec<String> {

82 pub fn provider_ids(&self) -> Vec<String> {

9083 let mut ids: Vec<_> = self

9184 .providers

9285 .keys()

Modified agent/llm/skill.rs +7 −33 View agent/llm/skill.rs

@@ -1,5 +1,3 @@

1//! Skill system — reusable prompt templates that agents can invoke.

2

31use crate::error::LlmError;

42use serde::{Deserialize, Serialize};

53use std::collections::HashMap;

@@ -23,23 +21,7 @@

2321}

2422

2523impl Skill {

26 /// Render the template by substituting `{{name}}` placeholders.

27 ///

28 /// Substitution is a **single left-to-right pass**: text written to the

29 /// output is never rescanned, so a parameter value that itself contains

30 /// `{{other}}` is emitted literally rather than expanded. This is what

31 /// stops caller-supplied data from being reinterpreted as template syntax,

32 /// and it makes the result independent of `params` iteration order.

33 ///

34 /// Resolution for a placeholder `{{k}}`:

35 /// 1. `k` is declared and present in `params` → that value

36 /// 2. `k` is declared with a `default` → the default

37 /// 3. `k` is declared and optional → empty string

38 /// 4. `k` is undeclared → emitted verbatim, including the braces

39 ///

40 /// Keys in `params` not declared on the skill are ignored: only the

41 /// declared parameter set is substitutable.

42 pub fn render(&self, params: &HashMap<String, String>) -> Result<String, LlmError> {

24 pub fn render(&self, params: &HashMap<String, String>) -> Result<String, LlmError> {

4325 for param in &self.parameters {

4426 if param.required && !params.contains_key(&param.name) && param.default.is_none() {

4527 return Err(LlmError::SkillMissingParameter {

@@ -55,7 +37,7 @@

5537 while let Some(open) = rest.find("{{") {

5638 let after_open = &rest[open + 2..];

5739 let Some(close) = after_open.find("}}") else {

58 break; // Unclosed brace: the remainder is literal text.

40 break;

5941 };

6042 let end = open + 2 + close + 2;

6143

@@ -96,31 +78,25 @@

9678 }

9779 }

9880

99 /// Register a skill, returning the previous skill of the same name if one

100 /// was displaced.

101 pub fn register(&mut self, skill: Skill) -> Option<Skill> {

81 pub fn register(&mut self, skill: Skill) -> Option<Skill> {

10282 self.skills.insert(skill.name.clone(), skill)

10383 }

10484

105 /// Get a skill by name.

106 pub fn get(&self, name: &str) -> Option<&Skill> {

85 pub fn get(&self, name: &str) -> Option<&Skill> {

10786 self.skills.get(name)

10887 }

10988

110 /// Look up and render a skill in one step.

111 pub fn render(&self, name: &str, params: &HashMap<String, String>) -> Result<String, LlmError> {

89 pub fn render(&self, name: &str, params: &HashMap<String, String>) -> Result<String, LlmError> {

11290 self.get(name)

11391 .ok_or_else(|| LlmError::SkillNotFound(name.to_string()))?

11492 .render(params)

11593 }

11694

117 /// List all skill names.

118 pub fn names(&self) -> Vec<String> {

95 pub fn names(&self) -> Vec<String> {

11996 self.skills.keys().cloned().collect()

12097 }

12198

122 /// Number of registered skills.

123 pub fn count(&self) -> usize {

99 pub fn count(&self) -> usize {

124100 self.skills.len()

125101 }

126102

@@ -168,8 +144,6 @@

168144 ("a".to_string(), "{{b}}".to_string()),

169145 ("b".to_string(), "SECRET".to_string()),

170146 ]);

171 // The `{{b}}` arriving as a's *value* must survive verbatim; the old

172 // two-pass implementation expanded it to "SECRET|SECRET".

173147 assert_eq!(s.render(&params).unwrap(), "{{b}}|SECRET");

174148 }

175149

Modified agent/llm/token.rs +10 −34 View agent/llm/token.rs

@@ -38,33 +38,14 @@

3838 }

3939}

4040

41/// Average characters per token for Latin-script text.

4241const LATIN_CHARS_PER_TOKEN: usize = 4;

4342

44/// Tokens for a low-detail image (OpenAI's documented flat cost).

4543pub const IMAGE_TOKENS_LOW: usize = 85;

4644

47/// Conservative upper estimate for a high-detail image (base + 4 tiles).

48///

49/// The real cost depends on pixel dimensions, which are not knowable from a

50/// URL. Over-estimating is the safe direction: it trips a local budget check

51/// rather than letting the request fail at the provider with a hard context

52/// overflow.

5345pub const IMAGE_TOKENS_HIGH: usize = 765;

5446

55/// Rough per-message framing overhead (role markers, separators).

5647pub const MESSAGE_OVERHEAD_TOKENS: usize = 4;

5748

58/// Estimate the token count of a string.

59///

60/// CJK characters count as roughly one token each; other scripts at

61/// [`LATIN_CHARS_PER_TOKEN`] characters per token. This counts **characters,

62/// not bytes** — counting UTF-8 bytes over-estimates CJK text by ~3x, which

63/// causes spurious budget rejections for Chinese/Japanese/Korean input.

64///

65/// This is a heuristic for budgeting and logging only. It is **not** accurate

66/// enough for billing and must not be the sole gate on a hard context limit;

67/// prefer a provider-supplied tokenizer wherever one is available.

6849pub fn estimate_text_tokens(text: &str) -> usize {

6950 let (cjk, other) = text.chars().fold((0usize, 0usize), |(cjk, other), c| {

7051 if is_cjk(c) {

@@ -76,25 +57,23 @@

7657 cjk + other.div_ceil(LATIN_CHARS_PER_TOKEN)

7758}

7859

79/// Estimate tokens for a JSON value (tool schemas, tool-call arguments).

8060pub fn estimate_json_tokens(value: &serde_json::Value) -> usize {

8161 serde_json::to_string(value)

8262 .map(|s| estimate_text_tokens(&s))

8363 .unwrap_or(0)

8464}

8565

86/// Characters that are typically one token each rather than ~4.

8766fn is_cjk(c: char) -> bool {

8867 matches!(c as u32,

89 0x2E80..=0x303F // CJK radicals and punctuation

90 | 0x3040..=0x30FF // Hiragana, Katakana

91 | 0x3400..=0x4DBF // CJK Unified Ext A

92 | 0x4E00..=0x9FFF // CJK Unified

93 | 0xA960..=0xA97F // Hangul Jamo Ext A

94 | 0xAC00..=0xD7FF // Hangul syllables

95 | 0xF900..=0xFAFF // CJK compatibility

96 | 0xFF00..=0xFFEF // Halfwidth and fullwidth forms

97 | 0x20000..=0x3FFFF // CJK Unified Ext B and beyond

68 0x2E80..=0x303F

69 | 0x3040..=0x30FF

70 | 0x3400..=0x4DBF

71 | 0x4E00..=0x9FFF

72 | 0xA960..=0xA97F

73 | 0xAC00..=0xD7FF

74 | 0xF900..=0xFAFF

75 | 0xFF00..=0xFFEF

76 | 0x20000..=0x3FFFF

9877 )

9978}

10079

@@ -138,8 +117,6 @@

138117

139118 #[test]

140119 fn cjk_counts_by_character_not_byte() {

141 // 12 chars, 36 UTF-8 bytes. Byte-based `len() / 4` treats this as 9

142 // tokens' worth of budget for text that really costs ~12.

143120 let zh = "这是一段中文测试文本内容";

144121 assert_eq!(zh.chars().count(), 12);

145122 assert_eq!(zh.len(), 36);

@@ -149,13 +126,12 @@

149126 #[test]

150127 fn latin_uses_four_chars_per_token() {

151128 assert_eq!(estimate_text_tokens("abcdefgh"), 2);

152 assert_eq!(estimate_text_tokens("abcde"), 2); // rounds up

129 assert_eq!(estimate_text_tokens("abcde"), 2);

153130 assert_eq!(estimate_text_tokens(""), 0);

154131 }

155132

156133 #[test]

157134 fn mixed_script_counts_each_run_appropriately() {

158 // 4 latin chars (1 token) + 2 CJK (2 tokens)

159135 assert_eq!(estimate_text_tokens("abcd中文"), 3);

160136 }

161137}