1use super::{Capability, CapabilityLocalization, CapabilityStatus};
20use crate::egress::{EgressError, EgressRequest, EgressRequestKind};
21use crate::tools::{Tool, ToolExecutionResult};
22use crate::traits::ToolContext;
23use async_trait::async_trait;
24use serde::{Deserialize, Serialize};
25use serde_json::{Value, json};
26
27pub const OPENROUTER_WORKSPACE_CAPABILITY_ID: &str = "openrouter_workspace";
28
29pub struct OpenRouterWorkspaceCapability;
32
33impl Capability for OpenRouterWorkspaceCapability {
34 fn id(&self) -> &str {
35 OPENROUTER_WORKSPACE_CAPABILITY_ID
36 }
37
38 fn name(&self) -> &str {
39 "OpenRouter Workspace"
40 }
41
42 fn description(&self) -> &str {
43 "Inspect OpenRouter workspace policy (budget, rate limits, tier) and detect incompatibilities with local routing configuration."
44 }
45
46 fn localizations(&self) -> Vec<CapabilityLocalization> {
47 vec![]
48 }
49
50 fn status(&self) -> CapabilityStatus {
51 CapabilityStatus::Available
52 }
53
54 fn icon(&self) -> Option<&str> {
55 Some("shield-check")
56 }
57
58 fn category(&self) -> Option<&str> {
59 Some("AI")
60 }
61
62 fn tools(&self) -> Vec<Box<dyn Tool>> {
63 vec![
64 Box::new(InspectOpenRouterWorkspaceTool),
65 Box::new(CheckOpenRouterPolicyCompatibilityTool),
66 ]
67 }
68}
69
70#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
79pub struct OpenRouterKeyInfo {
80 pub label: String,
82 pub usage_usd: f64,
84 pub limit_usd: Option<f64>,
86 pub remaining_usd: Option<f64>,
88 pub is_free_tier: bool,
90 pub rate_limit: Option<OpenRouterRateLimit>,
92}
93
94impl OpenRouterKeyInfo {
95 fn from_api_response(data: &Value) -> Result<Self, String> {
100 let label = data
101 .get("label")
102 .and_then(|v| v.as_str())
103 .unwrap_or("(unnamed)")
104 .to_string();
105
106 let usage_usd = data
108 .get("usage")
109 .and_then(|v| v.as_f64())
110 .map(|u| u / 1_000_000.0)
111 .unwrap_or(0.0);
112
113 let limit_usd = data
114 .get("limit")
115 .and_then(|v| if v.is_null() { None } else { v.as_f64() })
116 .map(|l| l / 1_000_000.0);
117
118 let remaining_usd = limit_usd.map(|l| (l - usage_usd).max(0.0));
119
120 let is_free_tier = data
121 .get("is_free_tier")
122 .and_then(|v| v.as_bool())
123 .unwrap_or(false);
124
125 let rate_limit = data.get("rate_limit").and_then(|rl| {
126 let requests = rl.get("requests")?.as_u64()?.try_into().ok()?;
127 let interval = rl.get("interval")?.as_str()?.to_string();
128 Some(OpenRouterRateLimit { requests, interval })
129 });
130
131 Ok(OpenRouterKeyInfo {
132 label,
133 usage_usd,
134 limit_usd,
135 remaining_usd,
136 is_free_tier,
137 rate_limit,
138 })
139 }
140}
141
142#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
144pub struct OpenRouterRateLimit {
145 pub requests: u32,
147 pub interval: String,
149}
150
151#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
154#[serde(tag = "kind", rename_all = "snake_case")]
155pub enum WorkspacePolicyDrift {
156 BudgetExhausted {
158 usage_usd: f64,
159 limit_usd: f64,
160 message: String,
161 },
162 BudgetBelowThreshold {
164 remaining_usd: f64,
165 threshold_usd: f64,
166 message: String,
167 },
168 FreeTierRestriction { feature: String, message: String },
171}
172
173#[derive(Debug, Clone, Serialize, Deserialize)]
175pub struct PolicyCompatibilityReport {
176 pub compatible: bool,
178 pub workspace: OpenRouterKeyInfo,
180 pub drifts: Vec<WorkspacePolicyDrift>,
182}
183
184pub fn detect_policy_drift(
187 workspace: &OpenRouterKeyInfo,
188 min_remaining_budget_usd: Option<f64>,
189 requires_paid_features: bool,
190) -> Vec<WorkspacePolicyDrift> {
191 let mut drifts = Vec::new();
192
193 if let (Some(remaining), Some(limit)) = (workspace.remaining_usd, workspace.limit_usd) {
195 if remaining <= 0.0 {
196 drifts.push(WorkspacePolicyDrift::BudgetExhausted {
197 usage_usd: workspace.usage_usd,
198 limit_usd: limit,
199 message: format!(
200 "Workspace budget exhausted: spent ${:.6} of ${:.6} limit",
201 workspace.usage_usd, limit
202 ),
203 });
204 } else if let Some(threshold) = min_remaining_budget_usd.filter(|&t| remaining < t) {
205 drifts.push(WorkspacePolicyDrift::BudgetBelowThreshold {
206 remaining_usd: remaining,
207 threshold_usd: threshold,
208 message: format!(
209 "Remaining workspace budget (${remaining:.6}) is below the requested \
210 threshold (${threshold:.6})"
211 ),
212 });
213 }
214 }
215 if workspace.is_free_tier && requires_paid_features {
219 drifts.push(WorkspacePolicyDrift::FreeTierRestriction {
220 feature: "paid_tier_routing".to_string(),
221 message: "Workspace is on the free tier, which restricts access to paid-only \
222 providers, ZDR endpoints, and data-retention controls. Upgrade to a paid \
223 key to use these features."
224 .to_string(),
225 });
226 }
227
228 drifts
229}
230
231struct InspectOpenRouterWorkspaceTool;
236
237#[async_trait]
238impl Tool for InspectOpenRouterWorkspaceTool {
239 fn name(&self) -> &str {
240 "inspect_openrouter_workspace"
241 }
242
243 fn description(&self) -> &str {
244 "Fetch OpenRouter workspace/key metadata (budget, rate limits, tier status). \
245 Returns structured policy metadata without exposing the API key itself. \
246 Use this to understand workspace constraints before configuring model routing."
247 }
248
249 fn parameters_schema(&self) -> Value {
250 json!({
251 "type": "object",
252 "properties": {},
253 "additionalProperties": false
254 })
255 }
256
257 async fn execute(&self, _arguments: Value) -> ToolExecutionResult {
258 ToolExecutionResult::tool_error(
259 "inspect_openrouter_workspace requires provider credentials — use execute_with_context",
260 )
261 }
262
263 async fn execute_with_context(
264 &self,
265 _arguments: Value,
266 context: &ToolContext,
267 ) -> ToolExecutionResult {
268 let body = match fetch_openrouter_key_info(context).await {
269 Ok(body) => body,
270 Err(e) => return ToolExecutionResult::tool_error(e),
271 };
272
273 let data = match body.get("data") {
274 Some(d) => d,
275 None => {
276 return ToolExecutionResult::tool_error(
277 "OpenRouter /auth/key response missing 'data' field",
278 );
279 }
280 };
281
282 let info = match OpenRouterKeyInfo::from_api_response(data) {
283 Ok(i) => i,
284 Err(e) => return ToolExecutionResult::tool_error(e),
285 };
286
287 ToolExecutionResult::Success(json!(info))
288 }
289}
290
291struct CheckOpenRouterPolicyCompatibilityTool;
296
297#[async_trait]
298impl Tool for CheckOpenRouterPolicyCompatibilityTool {
299 fn name(&self) -> &str {
300 "check_openrouter_policy_compatibility"
301 }
302
303 fn description(&self) -> &str {
304 "Check whether local routing configuration is compatible with the current OpenRouter \
305 workspace policy. Returns a compatibility report listing any detected drift between \
306 local settings and upstream workspace constraints (budget, tier, rate limits). \
307 Use before finalizing routing config to catch incompatibilities early."
308 }
309
310 fn parameters_schema(&self) -> Value {
311 json!({
312 "type": "object",
313 "properties": {
314 "min_remaining_budget_usd": {
315 "type": "number",
316 "minimum": 0,
317 "description": "Minimum remaining workspace budget (USD) required by your \
318 routing plan. A drift is reported if the workspace has less \
319 than this amount remaining. Omit to skip budget-threshold check."
320 },
321 "requires_paid_features": {
322 "type": "boolean",
323 "description": "Set true if the routing config uses paid-tier features such as \
324 ZDR endpoints, data-collection controls, or paid-only providers. \
325 A drift is reported when the workspace is on the free tier.",
326 "default": false
327 }
328 },
329 "additionalProperties": false
330 })
331 }
332
333 async fn execute(&self, _arguments: Value) -> ToolExecutionResult {
334 ToolExecutionResult::tool_error(
335 "check_openrouter_policy_compatibility requires provider credentials — use execute_with_context",
336 )
337 }
338
339 async fn execute_with_context(
340 &self,
341 arguments: Value,
342 context: &ToolContext,
343 ) -> ToolExecutionResult {
344 let min_remaining_budget_usd = arguments
345 .get("min_remaining_budget_usd")
346 .and_then(|v| v.as_f64())
347 .filter(|&v| v >= 0.0);
348 let requires_paid_features = arguments
349 .get("requires_paid_features")
350 .and_then(|v| v.as_bool())
351 .unwrap_or(false);
352
353 let body = match fetch_openrouter_key_info(context).await {
354 Ok(body) => body,
355 Err(e) => return ToolExecutionResult::tool_error(e),
356 };
357
358 let data = match body.get("data") {
359 Some(d) => d,
360 None => {
361 return ToolExecutionResult::tool_error(
362 "OpenRouter /auth/key response missing 'data' field",
363 );
364 }
365 };
366
367 let workspace = match OpenRouterKeyInfo::from_api_response(data) {
368 Ok(i) => i,
369 Err(e) => return ToolExecutionResult::tool_error(e),
370 };
371
372 let drifts =
373 detect_policy_drift(&workspace, min_remaining_budget_usd, requires_paid_features);
374 let compatible = drifts.is_empty();
375
376 let report = PolicyCompatibilityReport {
377 compatible,
378 workspace,
379 drifts,
380 };
381
382 ToolExecutionResult::Success(json!(report))
383 }
384}
385
386const OPENROUTER_KEY_INFO_URL: &str = "https://openrouter.ai/api/v1/auth/key";
391
392async fn fetch_openrouter_key_info(context: &ToolContext) -> Result<Value, String> {
393 let api_key = resolve_openrouter_key(context).await?;
394 let egress = context
395 .egress_service
396 .as_ref()
397 .ok_or_else(|| "OpenRouter workspace API requires the host egress service".to_string())?;
398
399 let response = egress
400 .send(
401 EgressRequest::new(
402 "GET",
403 OPENROUTER_KEY_INFO_URL,
404 EgressRequestKind::Capability,
405 )
406 .header("authorization", format!("Bearer {api_key}"))
407 .network_access(context.network_access.clone())
408 .timeout_ms(15_000),
409 )
410 .await
411 .map_err(|e| match e {
412 EgressError::NetworkAccessDenied { .. } => {
416 format!("OpenRouter workspace API blocked by network access policy: {e}")
417 }
418 other => format!("Failed to reach OpenRouter workspace API: {other}"),
419 })?;
420
421 if !(200..300).contains(&response.status) {
422 let body = String::from_utf8_lossy(&response.body);
423 return Err(format!(
424 "OpenRouter /auth/key returned HTTP {}: {}",
425 response.status, body
426 ));
427 }
428
429 serde_json::from_slice(&response.body)
430 .map_err(|e| format!("Failed to parse OpenRouter workspace response: {e}"))
431}
432
433async fn resolve_openrouter_key(context: &ToolContext) -> Result<String, String> {
434 let store = context
435 .provider_credential_store
436 .as_ref()
437 .ok_or_else(|| "No provider credential store available".to_string())?;
438
439 let creds = store
440 .get_default_provider_credentials("openrouter")
441 .await
442 .map_err(|e| format!("Failed to resolve OpenRouter credentials: {e}"))?
443 .ok_or_else(|| {
444 "No OpenRouter provider configured — add an OpenRouter provider to your org".to_string()
445 })?;
446
447 Ok(creds.api_key)
448}
449
450#[cfg(test)]
455mod tests {
456 use super::*;
457 use crate::egress::{EgressResponse, EgressService, EgressStreamResponse};
458 use crate::error::Result as CoreResult;
459 use crate::network_access::NetworkAccessList;
460 use crate::traits::{ProviderCredentialStore, ProviderCredentials};
461 use std::sync::{Arc, Mutex};
462
463 struct StaticCredentialStore;
464
465 #[async_trait]
466 impl ProviderCredentialStore for StaticCredentialStore {
467 async fn get_default_provider_credentials(
468 &self,
469 provider_type: &str,
470 ) -> CoreResult<Option<ProviderCredentials>> {
471 assert_eq!(provider_type, "openrouter");
472 Ok(Some(ProviderCredentials {
473 api_key: "test-key".to_string(),
474 base_url: None,
475 }))
476 }
477 }
478
479 #[derive(Default)]
480 struct RecordingEgress {
481 requests: Mutex<Vec<EgressRequest>>,
482 }
483
484 #[async_trait]
485 impl EgressService for RecordingEgress {
486 async fn send(
487 &self,
488 request: EgressRequest,
489 ) -> crate::egress::EgressResult<EgressResponse> {
490 self.requests.lock().unwrap().push(request);
491 Ok(EgressResponse {
492 status: 200,
493 headers: Default::default(),
494 body: br#"{"data":{"label":"egress-key","usage":1000000,"limit":2000000,"is_free_tier":false}}"#.to_vec(),
495 })
496 }
497
498 async fn send_stream(
499 &self,
500 _request: EgressRequest,
501 ) -> crate::egress::EgressResult<EgressStreamResponse> {
502 unreachable!("OpenRouter workspace tools use non-streaming egress")
503 }
504 }
505
506 fn unlimited_key() -> OpenRouterKeyInfo {
507 OpenRouterKeyInfo {
508 label: "test-key".to_string(),
509 usage_usd: 0.05,
510 limit_usd: None,
511 remaining_usd: None,
512 is_free_tier: false,
513 rate_limit: Some(OpenRouterRateLimit {
514 requests: 200,
515 interval: "10s".to_string(),
516 }),
517 }
518 }
519
520 fn capped_key(usage: f64, limit: f64) -> OpenRouterKeyInfo {
521 OpenRouterKeyInfo {
522 label: "capped-key".to_string(),
523 usage_usd: usage,
524 limit_usd: Some(limit),
525 remaining_usd: Some((limit - usage).max(0.0)),
526 is_free_tier: false,
527 rate_limit: None,
528 }
529 }
530
531 fn free_tier_key() -> OpenRouterKeyInfo {
532 OpenRouterKeyInfo {
533 label: "free-key".to_string(),
534 usage_usd: 0.0,
535 limit_usd: None,
536 remaining_usd: None,
537 is_free_tier: true,
538 rate_limit: Some(OpenRouterRateLimit {
539 requests: 20,
540 interval: "1m".to_string(),
541 }),
542 }
543 }
544
545 #[tokio::test]
546 async fn fetch_openrouter_key_info_uses_context_egress_service() {
547 let egress = Arc::new(RecordingEgress::default());
548 let network_access = NetworkAccessList::allow_only(["openrouter.ai"]);
552 assert!(
556 network_access.is_url_allowed(OPENROUTER_KEY_INFO_URL),
557 "network access list must allow the OpenRouter key-info URL"
558 );
559 let context = ToolContext::new(crate::SessionId::new())
560 .with_provider_credential_store(Arc::new(StaticCredentialStore))
561 .with_egress_service(egress.clone())
562 .with_network_access(Some(network_access.clone()));
563
564 let body = fetch_openrouter_key_info(&context).await.unwrap();
565
566 assert_eq!(body["data"]["label"], "egress-key");
567 let requests = egress.requests.lock().unwrap();
568 assert_eq!(requests.len(), 1);
569 let request = &requests[0];
570 assert_eq!(request.method, "GET");
571 assert_eq!(request.url, OPENROUTER_KEY_INFO_URL);
572 assert_eq!(request.kind, EgressRequestKind::Capability);
573 assert_eq!(
574 request.headers.get("authorization").map(String::as_str),
575 Some("Bearer test-key")
576 );
577 assert_eq!(request.network_access, Some(network_access));
578 assert_eq!(request.timeout_ms, Some(15_000));
579 }
580
581 #[test]
586 fn parse_unlimited_key_response() {
587 let raw = json!({
588 "label": "my-key",
589 "usage": 12_345,
590 "limit": null,
591 "is_free_tier": false,
592 "rate_limit": { "requests": 200, "interval": "10s" }
593 });
594 let info = OpenRouterKeyInfo::from_api_response(&raw).unwrap();
595 assert_eq!(info.label, "my-key");
596 assert!((info.usage_usd - 0.012345).abs() < 1e-9);
597 assert!(info.limit_usd.is_none());
598 assert!(info.remaining_usd.is_none());
599 assert!(!info.is_free_tier);
600 assert_eq!(info.rate_limit.as_ref().unwrap().requests, 200);
601 assert_eq!(info.rate_limit.as_ref().unwrap().interval, "10s");
602 }
603
604 #[test]
605 fn parse_capped_key_response() {
606 let raw = json!({
607 "label": "capped",
608 "usage": 500_000,
609 "limit": 1_000_000,
610 "is_free_tier": false
611 });
612 let info = OpenRouterKeyInfo::from_api_response(&raw).unwrap();
613 assert!((info.usage_usd - 0.5).abs() < 1e-9);
614 assert!((info.limit_usd.unwrap() - 1.0).abs() < 1e-9);
615 assert!((info.remaining_usd.unwrap() - 0.5).abs() < 1e-9);
616 }
617
618 #[test]
619 fn parse_exhausted_key_response() {
620 let raw = json!({
621 "label": "empty",
622 "usage": 1_000_000,
623 "limit": 1_000_000,
624 "is_free_tier": false
625 });
626 let info = OpenRouterKeyInfo::from_api_response(&raw).unwrap();
627 assert!((info.remaining_usd.unwrap()).abs() < 1e-9);
628 }
629
630 #[test]
631 fn parse_free_tier_key_response() {
632 let raw = json!({
633 "usage": 0,
634 "limit": null,
635 "is_free_tier": true,
636 "rate_limit": { "requests": 20, "interval": "1m" }
637 });
638 let info = OpenRouterKeyInfo::from_api_response(&raw).unwrap();
639 assert!(info.is_free_tier);
640 assert_eq!(info.rate_limit.as_ref().unwrap().requests, 20);
641 }
642
643 #[test]
644 fn parse_missing_label_uses_default() {
645 let raw = json!({ "usage": 0, "is_free_tier": false });
646 let info = OpenRouterKeyInfo::from_api_response(&raw).unwrap();
647 assert_eq!(info.label, "(unnamed)");
648 }
649
650 #[test]
655 fn no_drift_unlimited_key_no_constraints() {
656 let drifts = detect_policy_drift(&unlimited_key(), None, false);
657 assert!(drifts.is_empty());
658 }
659
660 #[test]
661 fn no_drift_unlimited_key_with_threshold() {
662 let drifts = detect_policy_drift(&unlimited_key(), Some(100.0), false);
664 assert!(drifts.is_empty());
665 }
666
667 #[test]
668 fn drift_budget_exhausted() {
669 let key = capped_key(1.0, 1.0);
670 let drifts = detect_policy_drift(&key, None, false);
671 assert_eq!(drifts.len(), 1);
672 assert!(matches!(
673 &drifts[0],
674 WorkspacePolicyDrift::BudgetExhausted { .. }
675 ));
676 }
677
678 #[test]
679 fn drift_budget_below_threshold() {
680 let key = capped_key(0.95, 1.0); let drifts = detect_policy_drift(&key, Some(0.10), false);
682 assert_eq!(drifts.len(), 1);
683 if let WorkspacePolicyDrift::BudgetBelowThreshold {
684 remaining_usd,
685 threshold_usd,
686 ..
687 } = &drifts[0]
688 {
689 assert!((remaining_usd - 0.05).abs() < 1e-9);
690 assert!((threshold_usd - 0.10).abs() < 1e-9);
691 } else {
692 panic!("wrong drift kind");
693 }
694 }
695
696 #[test]
697 fn no_drift_budget_above_threshold() {
698 let key = capped_key(0.5, 1.0); let drifts = detect_policy_drift(&key, Some(0.10), false);
700 assert!(drifts.is_empty());
701 }
702
703 #[test]
704 fn drift_free_tier_requires_paid() {
705 let key = free_tier_key();
706 let drifts = detect_policy_drift(&key, None, true);
707 assert_eq!(drifts.len(), 1);
708 assert!(matches!(
709 &drifts[0],
710 WorkspacePolicyDrift::FreeTierRestriction { .. }
711 ));
712 }
713
714 #[test]
715 fn no_drift_free_tier_no_paid_requirement() {
716 let key = free_tier_key();
717 let drifts = detect_policy_drift(&key, None, false);
718 assert!(drifts.is_empty());
719 }
720
721 #[test]
722 fn multiple_drifts_exhausted_and_free_tier() {
723 let mut key = capped_key(1.0, 1.0);
724 key.is_free_tier = true;
725 let drifts = detect_policy_drift(&key, None, true);
726 assert_eq!(drifts.len(), 2);
727 assert!(
728 drifts
729 .iter()
730 .any(|d| matches!(d, WorkspacePolicyDrift::BudgetExhausted { .. }))
731 );
732 assert!(
733 drifts
734 .iter()
735 .any(|d| matches!(d, WorkspacePolicyDrift::FreeTierRestriction { .. }))
736 );
737 }
738
739 #[test]
740 fn policy_compatibility_report_serializes() {
741 let report = PolicyCompatibilityReport {
742 compatible: false,
743 workspace: capped_key(1.0, 1.0),
744 drifts: vec![WorkspacePolicyDrift::BudgetExhausted {
745 usage_usd: 1.0,
746 limit_usd: 1.0,
747 message: "exhausted".to_string(),
748 }],
749 };
750 let json = serde_json::to_value(&report).unwrap();
751 assert_eq!(json["compatible"], false);
752 assert_eq!(json["drifts"][0]["kind"], "budget_exhausted");
753 }
754
755 #[test]
756 fn workspace_policy_drift_messages_are_non_empty() {
757 let key = capped_key(1.0, 1.0);
760 let drifts = detect_policy_drift(&key, Some(0.10), true);
761 for d in &drifts {
764 let msg = match d {
765 WorkspacePolicyDrift::BudgetExhausted { message, .. } => message.as_str(),
766 WorkspacePolicyDrift::BudgetBelowThreshold { message, .. } => message.as_str(),
767 WorkspacePolicyDrift::FreeTierRestriction { message, .. } => message.as_str(),
768 };
769 assert!(!msg.is_empty(), "drift message should not be empty");
770 }
771 }
772}