1use crate::{DrivenError, Result};
34use std::collections::HashMap;
35
36pub use dcp::binary::{
38 ArgType, BinaryMessageEnvelope, ChunkFlags, Flags, MessageType, SignedInvocation,
39 SignedToolDef, StreamChunk, ToolInvocation,
40};
41pub use dcp::capability::CapabilityManifest;
42pub use dcp::compat::adapter::AdapterError;
43pub use dcp::compat::json_rpc::RequestId;
44pub use dcp::compat::{JsonRpcError, JsonRpcParser, JsonRpcRequest, JsonRpcResponse, McpAdapter};
45pub use dcp::security::{Signer, Verifier};
46
47#[derive(Debug, Clone)]
52pub struct ZctiBuilder {
53 tool_id: u32,
55 arg_layout: u64,
57 args_data: Vec<u8>,
59 arg_index: usize,
61}
62
63impl ZctiBuilder {
64 pub fn new(tool_id: u32) -> Self {
66 Self {
67 tool_id,
68 arg_layout: 0,
69 args_data: Vec::new(),
70 arg_index: 0,
71 }
72 }
73
74 pub fn add_null(mut self) -> Self {
76 self.set_arg_type(ArgType::Null);
77 self
78 }
79
80 pub fn add_bool(mut self, value: bool) -> Self {
82 self.set_arg_type(ArgType::Bool);
83 self.args_data.push(if value { 1 } else { 0 });
84 self
85 }
86
87 pub fn add_i32(mut self, value: i32) -> Self {
89 self.set_arg_type(ArgType::I32);
90 self.args_data.extend_from_slice(&value.to_le_bytes());
91 self
92 }
93
94 pub fn add_i64(mut self, value: i64) -> Self {
96 self.set_arg_type(ArgType::I64);
97 self.args_data.extend_from_slice(&value.to_le_bytes());
98 self
99 }
100
101 pub fn add_f64(mut self, value: f64) -> Self {
103 self.set_arg_type(ArgType::F64);
104 self.args_data.extend_from_slice(&value.to_le_bytes());
105 self
106 }
107
108 pub fn add_string(mut self, value: &str) -> Self {
110 self.set_arg_type(ArgType::String);
111 let len = value.len() as u32;
113 self.args_data.extend_from_slice(&len.to_le_bytes());
114 self.args_data.extend_from_slice(value.as_bytes());
115 self
116 }
117
118 pub fn add_bytes(mut self, value: &[u8]) -> Self {
120 self.set_arg_type(ArgType::Bytes);
121 let len = value.len() as u32;
123 self.args_data.extend_from_slice(&len.to_le_bytes());
124 self.args_data.extend_from_slice(value);
125 self
126 }
127
128 pub fn add_raw(mut self, data: &[u8]) -> Self {
130 self.args_data.extend_from_slice(data);
131 self
132 }
133
134 fn set_arg_type(&mut self, arg_type: ArgType) {
136 if self.arg_index < ToolInvocation::MAX_ARGS {
137 let shift = self.arg_index * 4;
138 self.arg_layout |= (arg_type as u64) << shift;
139 self.arg_index += 1;
140 }
141 }
142
143 pub fn build(self) -> ToolInvocation {
145 ToolInvocation::new(
146 self.tool_id,
147 self.arg_layout,
148 0, self.args_data.len() as u32,
150 )
151 }
152
153 pub fn build_with_data(self) -> (ToolInvocation, Vec<u8>) {
155 let invocation = ToolInvocation::new(
156 self.tool_id,
157 self.arg_layout,
158 0,
159 self.args_data.len() as u32,
160 );
161 (invocation, self.args_data)
162 }
163
164 pub fn arg_count(&self) -> usize {
166 self.arg_index
167 }
168
169 pub fn data_size(&self) -> usize {
171 self.args_data.len()
172 }
173}
174
175#[derive(Debug)]
179pub struct ZctiReader<'a> {
180 invocation: ToolInvocation,
182 data: &'a [u8],
184 offset: usize,
186 arg_index: usize,
188}
189
190impl<'a> ZctiReader<'a> {
191 pub fn new(invocation: &ToolInvocation, data: &'a [u8]) -> Self {
193 Self {
194 invocation: *invocation,
195 data,
196 offset: 0,
197 arg_index: 0,
198 }
199 }
200
201 pub fn tool_id(&self) -> u32 {
203 self.invocation.tool_id
204 }
205
206 pub fn arg_count(&self) -> usize {
208 self.invocation.arg_count()
209 }
210
211 pub fn peek_type(&self) -> Option<ArgType> {
213 self.invocation.get_arg_type(self.arg_index)
214 }
215
216 pub fn read_bool(&mut self) -> Result<bool> {
218 self.check_type(ArgType::Bool)?;
219 if self.offset >= self.data.len() {
220 return Err(DrivenError::InvalidBinary(
221 "Insufficient data for bool".to_string(),
222 ));
223 }
224 let value = self.data[self.offset] != 0;
225 self.offset += 1;
226 self.arg_index += 1;
227 Ok(value)
228 }
229
230 pub fn read_i32(&mut self) -> Result<i32> {
232 self.check_type(ArgType::I32)?;
233 if self.offset + 4 > self.data.len() {
234 return Err(DrivenError::InvalidBinary(
235 "Insufficient data for i32".to_string(),
236 ));
237 }
238 let bytes: [u8; 4] = self.data[self.offset..self.offset + 4].try_into().unwrap();
239 let value = i32::from_le_bytes(bytes);
240 self.offset += 4;
241 self.arg_index += 1;
242 Ok(value)
243 }
244
245 pub fn read_i64(&mut self) -> Result<i64> {
247 self.check_type(ArgType::I64)?;
248 if self.offset + 8 > self.data.len() {
249 return Err(DrivenError::InvalidBinary(
250 "Insufficient data for i64".to_string(),
251 ));
252 }
253 let bytes: [u8; 8] = self.data[self.offset..self.offset + 8].try_into().unwrap();
254 let value = i64::from_le_bytes(bytes);
255 self.offset += 8;
256 self.arg_index += 1;
257 Ok(value)
258 }
259
260 pub fn read_f64(&mut self) -> Result<f64> {
262 self.check_type(ArgType::F64)?;
263 if self.offset + 8 > self.data.len() {
264 return Err(DrivenError::InvalidBinary(
265 "Insufficient data for f64".to_string(),
266 ));
267 }
268 let bytes: [u8; 8] = self.data[self.offset..self.offset + 8].try_into().unwrap();
269 let value = f64::from_le_bytes(bytes);
270 self.offset += 8;
271 self.arg_index += 1;
272 Ok(value)
273 }
274
275 pub fn read_string(&mut self) -> Result<&'a str> {
277 self.check_type(ArgType::String)?;
278 if self.offset + 4 > self.data.len() {
279 return Err(DrivenError::InvalidBinary(
280 "Insufficient data for string length".to_string(),
281 ));
282 }
283 let len_bytes: [u8; 4] = self.data[self.offset..self.offset + 4].try_into().unwrap();
284 let len = u32::from_le_bytes(len_bytes) as usize;
285 self.offset += 4;
286
287 if self.offset + len > self.data.len() {
288 return Err(DrivenError::InvalidBinary(
289 "Insufficient data for string content".to_string(),
290 ));
291 }
292 let value = std::str::from_utf8(&self.data[self.offset..self.offset + len])
293 .map_err(|e| DrivenError::InvalidBinary(format!("Invalid UTF-8: {}", e)))?;
294 self.offset += len;
295 self.arg_index += 1;
296 Ok(value)
297 }
298
299 pub fn read_bytes(&mut self) -> Result<&'a [u8]> {
301 self.check_type(ArgType::Bytes)?;
302 if self.offset + 4 > self.data.len() {
303 return Err(DrivenError::InvalidBinary(
304 "Insufficient data for bytes length".to_string(),
305 ));
306 }
307 let len_bytes: [u8; 4] = self.data[self.offset..self.offset + 4].try_into().unwrap();
308 let len = u32::from_le_bytes(len_bytes) as usize;
309 self.offset += 4;
310
311 if self.offset + len > self.data.len() {
312 return Err(DrivenError::InvalidBinary(
313 "Insufficient data for bytes content".to_string(),
314 ));
315 }
316 let value = &self.data[self.offset..self.offset + len];
317 self.offset += len;
318 self.arg_index += 1;
319 Ok(value)
320 }
321
322 pub fn skip(&mut self) -> Result<()> {
324 let arg_type = self
325 .peek_type()
326 .ok_or_else(|| DrivenError::InvalidBinary("No more arguments".to_string()))?;
327
328 match arg_type {
329 ArgType::Null => {
330 self.arg_index += 1;
331 }
332 ArgType::Bool => {
333 self.offset += 1;
334 self.arg_index += 1;
335 }
336 ArgType::I32 => {
337 self.offset += 4;
338 self.arg_index += 1;
339 }
340 ArgType::I64 | ArgType::F64 => {
341 self.offset += 8;
342 self.arg_index += 1;
343 }
344 ArgType::String | ArgType::Bytes => {
345 if self.offset + 4 > self.data.len() {
346 return Err(DrivenError::InvalidBinary(
347 "Insufficient data for length".to_string(),
348 ));
349 }
350 let len_bytes: [u8; 4] =
351 self.data[self.offset..self.offset + 4].try_into().unwrap();
352 let len = u32::from_le_bytes(len_bytes) as usize;
353 self.offset += 4 + len;
354 self.arg_index += 1;
355 }
356 ArgType::Array | ArgType::Object => {
357 return Err(DrivenError::InvalidBinary(
359 "Complex types not yet supported".to_string(),
360 ));
361 }
362 }
363 Ok(())
364 }
365
366 pub fn has_more(&self) -> bool {
368 self.arg_index < self.arg_count()
369 }
370
371 pub fn remaining_data(&self) -> &'a [u8] {
373 &self.data[self.offset..]
374 }
375
376 fn check_type(&self, expected: ArgType) -> Result<()> {
378 let actual = self
379 .peek_type()
380 .ok_or_else(|| DrivenError::InvalidBinary("No more arguments".to_string()))?;
381 if actual != expected {
382 return Err(DrivenError::InvalidBinary(format!(
383 "Type mismatch: expected {:?}, got {:?}",
384 expected, actual
385 )));
386 }
387 Ok(())
388 }
389}
390
391#[derive(Debug)]
396pub struct SharedArgBuffer {
397 data: Vec<u8>,
399 write_offset: usize,
401}
402
403impl SharedArgBuffer {
404 pub fn new(capacity: usize) -> Self {
406 Self {
407 data: vec![0u8; capacity],
408 write_offset: 0,
409 }
410 }
411
412 pub fn with_default_capacity() -> Self {
414 Self::new(65536)
415 }
416
417 pub fn write(&mut self, data: &[u8]) -> Result<u32> {
419 if self.write_offset + data.len() > self.data.len() {
420 return Err(DrivenError::InvalidBinary("Buffer overflow".to_string()));
421 }
422 let offset = self.write_offset as u32;
423 self.data[self.write_offset..self.write_offset + data.len()].copy_from_slice(data);
424 self.write_offset += data.len();
425 Ok(offset)
426 }
427
428 pub fn read(&self, offset: u32, len: u32) -> Result<&[u8]> {
430 let start = offset as usize;
431 let end = start + len as usize;
432 if end > self.data.len() {
433 return Err(DrivenError::InvalidBinary("Read out of bounds".to_string()));
434 }
435 Ok(&self.data[start..end])
436 }
437
438 pub fn offset(&self) -> u32 {
440 self.write_offset as u32
441 }
442
443 pub fn reset(&mut self) {
445 self.write_offset = 0;
446 }
447
448 pub fn data(&self) -> &[u8] {
450 &self.data[..self.write_offset]
451 }
452
453 pub fn capacity(&self) -> usize {
455 self.data.len()
456 }
457
458 pub fn remaining(&self) -> usize {
460 self.data.len() - self.write_offset
461 }
462}
463
464#[derive(Debug, Clone)]
466pub struct DcpConfig {
467 pub enabled: bool,
469 pub prefer_dcp: bool,
471 pub endpoint: Option<String>,
473 pub timeout_ms: u64,
475 pub signing_enabled: bool,
477 pub signing_seed: Option<[u8; 32]>,
479}
480
481impl Default for DcpConfig {
482 fn default() -> Self {
483 Self {
484 enabled: true,
485 prefer_dcp: true,
486 endpoint: None,
487 timeout_ms: 5000,
488 signing_enabled: false,
489 signing_seed: None,
490 }
491 }
492}
493
494#[derive(Debug, Clone, Copy, PartialEq, Eq)]
496pub enum ConnectionState {
497 Disconnected,
499 Connecting,
501 ConnectedDcp,
503 ConnectedMcp,
505 Failed,
507}
508
509#[derive(Debug, Clone, Copy, PartialEq, Eq)]
511pub enum Protocol {
512 Dcp,
514 Mcp,
516}
517
518#[derive(Debug)]
520pub enum InvocationResult {
521 Dcp(Vec<u8>),
523 Mcp(String),
525}
526
527impl InvocationResult {
528 pub fn is_dcp(&self) -> bool {
530 matches!(self, InvocationResult::Dcp(_))
531 }
532
533 pub fn is_mcp(&self) -> bool {
535 matches!(self, InvocationResult::Mcp(_))
536 }
537
538 pub fn as_dcp(&self) -> Option<&[u8]> {
540 match self {
541 InvocationResult::Dcp(bytes) => Some(bytes),
542 _ => None,
543 }
544 }
545
546 pub fn as_mcp(&self) -> Option<&str> {
548 match self {
549 InvocationResult::Mcp(json) => Some(json),
550 _ => None,
551 }
552 }
553}
554
555pub struct DcpClient {
560 config: DcpConfig,
562 state: ConnectionState,
564 mcp_adapter: Option<McpAdapter>,
566 capabilities: CapabilityManifest,
568 signer: Option<Signer>,
570 tools: HashMap<u32, ToolDefinition>,
572 next_tool_id: u32,
574}
575
576#[derive(Debug, Clone)]
578pub struct ToolDefinition {
579 pub id: u32,
581 pub name: String,
583 pub description: String,
585 pub schema_hash: [u8; 32],
587 pub capabilities: u64,
589}
590
591#[derive(Debug, Clone, Copy)]
593pub struct CapabilityStats {
594 pub tool_count: u32,
596 pub resource_count: u32,
598 pub prompt_count: u32,
600 pub extension_count: u32,
602 pub version: u16,
604}
605
606#[derive(Debug, Clone)]
608pub struct DcpMessage {
609 pub message_type: MessageType,
611 pub flags: u8,
613 pub payload: Vec<u8>,
615}
616
617impl DcpMessage {
618 pub fn new(message_type: MessageType, flags: u8, payload: Vec<u8>) -> Self {
620 Self {
621 message_type,
622 flags,
623 payload,
624 }
625 }
626
627 pub fn tool(payload: Vec<u8>) -> Self {
629 Self::new(MessageType::Tool, 0, payload)
630 }
631
632 pub fn resource(payload: Vec<u8>) -> Self {
634 Self::new(MessageType::Resource, 0, payload)
635 }
636
637 pub fn prompt(payload: Vec<u8>) -> Self {
639 Self::new(MessageType::Prompt, 0, payload)
640 }
641
642 pub fn response(payload: Vec<u8>) -> Self {
644 Self::new(MessageType::Response, 0, payload)
645 }
646
647 pub fn error(payload: Vec<u8>) -> Self {
649 Self::new(MessageType::Error, 0, payload)
650 }
651
652 pub fn stream(payload: Vec<u8>) -> Self {
654 Self::new(MessageType::Stream, Flags::STREAMING, payload)
655 }
656
657 pub fn is_streaming(&self) -> bool {
659 self.flags & Flags::STREAMING != 0
660 }
661
662 pub fn is_compressed(&self) -> bool {
664 self.flags & Flags::COMPRESSED != 0
665 }
666
667 pub fn is_signed(&self) -> bool {
669 self.flags & Flags::SIGNED != 0
670 }
671
672 pub fn with_streaming(mut self) -> Self {
674 self.flags |= Flags::STREAMING;
675 self
676 }
677
678 pub fn with_compressed(mut self) -> Self {
680 self.flags |= Flags::COMPRESSED;
681 self
682 }
683
684 pub fn with_signed(mut self) -> Self {
686 self.flags |= Flags::SIGNED;
687 self
688 }
689
690 pub fn encode(&self) -> Vec<u8> {
692 let envelope =
693 BinaryMessageEnvelope::new(self.message_type, self.flags, self.payload.len() as u32);
694
695 let mut result = Vec::with_capacity(BinaryMessageEnvelope::SIZE + self.payload.len());
696 result.extend_from_slice(envelope.as_bytes());
697 result.extend_from_slice(&self.payload);
698 result
699 }
700
701 pub fn decode(bytes: &[u8]) -> Result<Self> {
703 if bytes.len() < BinaryMessageEnvelope::SIZE {
704 return Err(DrivenError::InvalidBinary(
705 "Insufficient data for BME header".to_string(),
706 ));
707 }
708
709 let envelope = BinaryMessageEnvelope::from_bytes(bytes)
710 .map_err(|e| DrivenError::InvalidBinary(format!("Invalid BME: {:?}", e)))?;
711
712 let message_type = envelope.get_message_type().ok_or_else(|| {
713 DrivenError::InvalidBinary(format!("Unknown message type: {}", envelope.message_type))
714 })?;
715
716 let payload_len = envelope.payload_len as usize;
717 let total_len = BinaryMessageEnvelope::SIZE + payload_len;
718
719 if bytes.len() < total_len {
720 return Err(DrivenError::InvalidBinary(format!(
721 "Insufficient data: expected {} bytes, got {}",
722 total_len,
723 bytes.len()
724 )));
725 }
726
727 let payload = bytes[BinaryMessageEnvelope::SIZE..total_len].to_vec();
728
729 Ok(Self {
730 message_type,
731 flags: envelope.flags,
732 payload,
733 })
734 }
735}
736
737#[derive(Debug)]
739pub struct StreamAssembler {
740 chunks: Vec<StreamChunk>,
742 data: Vec<u8>,
744 next_sequence: u32,
746 complete: bool,
748 error: bool,
750}
751
752impl StreamAssembler {
753 pub fn new() -> Self {
755 Self {
756 chunks: Vec::new(),
757 data: Vec::new(),
758 next_sequence: 0,
759 complete: false,
760 error: false,
761 }
762 }
763
764 pub fn add_chunk(&mut self, chunk_bytes: &[u8], payload: &[u8]) -> Result<()> {
766 if self.complete {
767 return Err(DrivenError::InvalidBinary(
768 "Stream already complete".to_string(),
769 ));
770 }
771
772 let chunk = StreamChunk::from_bytes(chunk_bytes)
773 .map_err(|e| DrivenError::InvalidBinary(format!("Invalid chunk: {:?}", e)))?;
774
775 let sequence = chunk.sequence;
777 let flags = chunk.flags;
778 let len = chunk.len;
779
780 if sequence != self.next_sequence {
782 return Err(DrivenError::InvalidBinary(format!(
783 "Out of order chunk: expected {}, got {}",
784 self.next_sequence, sequence
785 )));
786 }
787
788 if flags & ChunkFlags::ERROR != 0 {
790 self.error = true;
791 self.complete = true;
792 return Err(DrivenError::InvalidBinary("Stream error".to_string()));
793 }
794
795 if payload.len() != len as usize {
797 return Err(DrivenError::InvalidBinary(format!(
798 "Payload length mismatch: header says {}, got {}",
799 len,
800 payload.len()
801 )));
802 }
803
804 self.data.extend_from_slice(payload);
806 self.chunks.push(*chunk);
807 self.next_sequence += 1;
808
809 if flags & ChunkFlags::LAST != 0 {
811 self.complete = true;
812 }
813
814 Ok(())
815 }
816
817 pub fn is_complete(&self) -> bool {
819 self.complete
820 }
821
822 pub fn has_error(&self) -> bool {
824 self.error
825 }
826
827 pub fn data(&self) -> Option<&[u8]> {
829 if self.complete && !self.error {
830 Some(&self.data)
831 } else {
832 None
833 }
834 }
835
836 pub fn take_data(self) -> Option<Vec<u8>> {
838 if self.complete && !self.error {
839 Some(self.data)
840 } else {
841 None
842 }
843 }
844
845 pub fn chunk_count(&self) -> usize {
847 self.chunks.len()
848 }
849}
850
851impl Default for StreamAssembler {
852 fn default() -> Self {
853 Self::new()
854 }
855}
856
857#[derive(Debug)]
859pub struct StreamBuilder {
860 chunk_size: usize,
862 sequence: u32,
864}
865
866impl StreamBuilder {
867 pub fn new(chunk_size: usize) -> Self {
869 Self {
870 chunk_size,
871 sequence: 0,
872 }
873 }
874
875 pub fn with_default_chunk_size() -> Self {
877 Self::new(65536)
878 }
879
880 pub fn build_chunks(&mut self, data: &[u8]) -> Vec<(StreamChunk, Vec<u8>)> {
882 let mut chunks = Vec::new();
883 let mut offset = 0;
884
885 while offset < data.len() {
886 let remaining = data.len() - offset;
887 let chunk_len = remaining.min(self.chunk_size);
888 let is_first = offset == 0;
889 let is_last = offset + chunk_len >= data.len();
890
891 let flags = if is_first && is_last {
892 ChunkFlags::FIRST | ChunkFlags::LAST
893 } else if is_first {
894 ChunkFlags::FIRST
895 } else if is_last {
896 ChunkFlags::LAST
897 } else {
898 ChunkFlags::CONTINUE
899 };
900
901 let chunk = StreamChunk::new(self.sequence, flags, chunk_len as u16);
902 let payload = data[offset..offset + chunk_len].to_vec();
903
904 chunks.push((chunk, payload));
905 self.sequence += 1;
906 offset += chunk_len;
907 }
908
909 chunks
910 }
911
912 pub fn reset(&mut self) {
914 self.sequence = 0;
915 }
916}
917
918impl DcpClient {
919 pub fn new(config: DcpConfig) -> Self {
921 let signer = if config.signing_enabled {
922 config.signing_seed.map(|seed| Signer::from_seed(&seed))
923 } else {
924 None
925 };
926
927 Self {
928 config,
929 state: ConnectionState::Disconnected,
930 mcp_adapter: None,
931 capabilities: CapabilityManifest::new(1),
932 signer,
933 tools: HashMap::new(),
934 next_tool_id: 1,
935 }
936 }
937
938 pub fn with_defaults() -> Self {
940 Self::new(DcpConfig::default())
941 }
942
943 pub fn state(&self) -> ConnectionState {
945 self.state
946 }
947
948 pub fn is_connected(&self) -> bool {
950 matches!(
951 self.state,
952 ConnectionState::ConnectedDcp | ConnectionState::ConnectedMcp
953 )
954 }
955
956 pub fn is_dcp_connected(&self) -> bool {
958 self.state == ConnectionState::ConnectedDcp
959 }
960
961 pub fn is_mcp_connected(&self) -> bool {
963 self.state == ConnectionState::ConnectedMcp
964 }
965
966 pub fn prefer_dcp(&self) -> bool {
968 self.config.prefer_dcp && self.config.enabled
969 }
970
971 pub fn set_prefer_dcp(&mut self, prefer: bool) {
973 self.config.prefer_dcp = prefer;
974 }
975
976 pub fn enable_dcp(&mut self) {
978 self.config.enabled = true;
979 }
980
981 pub fn disable_dcp(&mut self) {
983 self.config.enabled = false;
984 }
985
986 pub fn is_dcp_enabled(&self) -> bool {
988 self.config.enabled
989 }
990
991 pub fn current_protocol(&self) -> Option<Protocol> {
993 match self.state {
994 ConnectionState::ConnectedDcp => Some(Protocol::Dcp),
995 ConnectionState::ConnectedMcp => Some(Protocol::Mcp),
996 _ => None,
997 }
998 }
999
1000 pub fn select_protocol(&self) -> Option<Protocol> {
1005 if self.is_dcp_connected() && self.prefer_dcp() {
1006 Some(Protocol::Dcp)
1007 } else if self.is_mcp_connected() || self.has_mcp_fallback() {
1008 Some(Protocol::Mcp)
1009 } else if self.is_dcp_connected() {
1010 Some(Protocol::Dcp)
1011 } else {
1012 None
1013 }
1014 }
1015
1016 pub fn invoke_tool_auto(&self, tool_id: u32, args: &[u8]) -> Result<InvocationResult> {
1020 match self.select_protocol() {
1021 Some(Protocol::Dcp) => {
1022 let encoded = self.invoke_tool(tool_id, args)?;
1023 Ok(InvocationResult::Dcp(encoded))
1024 }
1025 Some(Protocol::Mcp) => {
1026 let tool = self
1028 .tools
1029 .get(&tool_id)
1030 .ok_or_else(|| DrivenError::Config(format!("Tool {} not found", tool_id)))?;
1031
1032 let request = serde_json::json!({
1033 "jsonrpc": "2.0",
1034 "method": "tools/call",
1035 "params": {
1036 "name": tool.name,
1037 "arguments": serde_json::from_slice::<serde_json::Value>(args)
1038 .unwrap_or(serde_json::Value::Null)
1039 },
1040 "id": 1
1041 });
1042
1043 let response = self.handle_mcp_request(&request.to_string())?;
1044 Ok(InvocationResult::Mcp(response))
1045 }
1046 None => Err(DrivenError::Config("No protocol available".to_string())),
1047 }
1048 }
1049
1050 pub fn capabilities(&self) -> &CapabilityManifest {
1052 &self.capabilities
1053 }
1054
1055 pub fn capabilities_mut(&mut self) -> &mut CapabilityManifest {
1057 &mut self.capabilities
1058 }
1059
1060 pub fn connect(&mut self, endpoint: &str) -> Result<()> {
1064 self.state = ConnectionState::Connecting;
1065
1066 if self.config.enabled {
1068 match self.try_dcp_connect(endpoint) {
1069 Ok(()) => {
1070 self.state = ConnectionState::ConnectedDcp;
1071 return Ok(());
1072 }
1073 Err(e) => {
1074 tracing::warn!("DCP connection failed, falling back to MCP: {}", e);
1076 }
1077 }
1078 }
1079
1080 match self.try_mcp_connect(endpoint) {
1082 Ok(()) => {
1083 self.state = ConnectionState::ConnectedMcp;
1084 Ok(())
1085 }
1086 Err(e) => {
1087 self.state = ConnectionState::Failed;
1088 Err(e)
1089 }
1090 }
1091 }
1092
1093 fn try_dcp_connect(&mut self, _endpoint: &str) -> Result<()> {
1095 Err(DrivenError::Config(
1099 "DCP connection not yet implemented".to_string(),
1100 ))
1101 }
1102
1103 fn try_mcp_connect(&mut self, _endpoint: &str) -> Result<()> {
1105 self.mcp_adapter = Some(McpAdapter::new());
1107 Ok(())
1108 }
1109
1110 pub fn disconnect(&mut self) {
1112 self.state = ConnectionState::Disconnected;
1113 self.mcp_adapter = None;
1114 }
1115
1116 pub fn register_tool(
1118 &mut self,
1119 name: &str,
1120 description: &str,
1121 schema_hash: [u8; 32],
1122 capabilities: u64,
1123 ) -> u32 {
1124 let id = self.next_tool_id;
1125 self.next_tool_id += 1;
1126
1127 let tool = ToolDefinition {
1128 id,
1129 name: name.to_string(),
1130 description: description.to_string(),
1131 schema_hash,
1132 capabilities,
1133 };
1134
1135 self.tools.insert(id, tool);
1136
1137 if id < CapabilityManifest::MAX_TOOLS as u32 {
1139 self.capabilities.set_tool(id as u16);
1140 }
1141
1142 id
1143 }
1144
1145 pub fn get_tool(&self, id: u32) -> Option<&ToolDefinition> {
1147 self.tools.get(&id)
1148 }
1149
1150 pub fn list_tools(&self) -> Vec<&ToolDefinition> {
1152 self.tools.values().collect()
1153 }
1154
1155 pub fn invoke_tool(&self, tool_id: u32, args: &[u8]) -> Result<Vec<u8>> {
1159 if !self.is_connected() {
1160 return Err(DrivenError::Config("Not connected".to_string()));
1161 }
1162
1163 if !self.tools.contains_key(&tool_id) {
1165 return Err(DrivenError::Config(format!("Tool {} not found", tool_id)));
1166 }
1167
1168 let invocation = ToolInvocation::new(tool_id, 0, 0, args.len() as u32);
1170
1171 let mut payload = Vec::with_capacity(ToolInvocation::SIZE + args.len());
1173 payload.extend_from_slice(&invocation.as_bytes());
1174 payload.extend_from_slice(args);
1175
1176 let flags = if self.signer.is_some() {
1178 Flags::SIGNED
1179 } else {
1180 0
1181 };
1182 let message = DcpMessage::new(MessageType::Tool, flags, payload);
1183
1184 let encoded = message.encode();
1186
1187 Ok(encoded)
1190 }
1191
1192 pub fn encode_message(&self, message_type: MessageType, payload: &[u8]) -> Vec<u8> {
1194 let flags = if self.signer.is_some() {
1195 Flags::SIGNED
1196 } else {
1197 0
1198 };
1199 let message = DcpMessage::new(message_type, flags, payload.to_vec());
1200 message.encode()
1201 }
1202
1203 pub fn decode_message(&self, bytes: &[u8]) -> Result<DcpMessage> {
1205 DcpMessage::decode(bytes)
1206 }
1207
1208 pub fn create_tool_message(&self, tool_id: u32, args: &[u8]) -> Result<DcpMessage> {
1210 if !self.tools.contains_key(&tool_id) {
1211 return Err(DrivenError::Config(format!("Tool {} not found", tool_id)));
1212 }
1213
1214 let invocation = ToolInvocation::new(tool_id, 0, 0, args.len() as u32);
1215
1216 let mut payload = Vec::with_capacity(ToolInvocation::SIZE + args.len());
1217 payload.extend_from_slice(&invocation.as_bytes());
1218 payload.extend_from_slice(args);
1219
1220 let flags = if self.signer.is_some() {
1221 Flags::SIGNED
1222 } else {
1223 0
1224 };
1225 Ok(DcpMessage::new(MessageType::Tool, flags, payload))
1226 }
1227
1228 pub fn create_resource_message(&self, resource_uri: &str) -> DcpMessage {
1230 let payload = resource_uri.as_bytes().to_vec();
1231 let flags = if self.signer.is_some() {
1232 Flags::SIGNED
1233 } else {
1234 0
1235 };
1236 DcpMessage::new(MessageType::Resource, flags, payload)
1237 }
1238
1239 pub fn create_prompt_message(&self, prompt: &str) -> DcpMessage {
1241 let payload = prompt.as_bytes().to_vec();
1242 let flags = if self.signer.is_some() {
1243 Flags::SIGNED
1244 } else {
1245 0
1246 };
1247 DcpMessage::new(MessageType::Prompt, flags, payload)
1248 }
1249
1250 pub fn create_response_message(&self, response: &[u8]) -> DcpMessage {
1252 let flags = if self.signer.is_some() {
1253 Flags::SIGNED
1254 } else {
1255 0
1256 };
1257 DcpMessage::new(MessageType::Response, flags, response.to_vec())
1258 }
1259
1260 pub fn create_error_message(&self, error_code: u32, error_msg: &str) -> DcpMessage {
1262 let mut payload = Vec::with_capacity(4 + error_msg.len());
1263 payload.extend_from_slice(&error_code.to_le_bytes());
1264 payload.extend_from_slice(error_msg.as_bytes());
1265 DcpMessage::new(MessageType::Error, 0, payload)
1266 }
1267
1268 pub fn create_stream_builder(&self, chunk_size: usize) -> StreamBuilder {
1270 StreamBuilder::new(chunk_size)
1271 }
1272
1273 pub fn create_stream_assembler(&self) -> StreamAssembler {
1275 StreamAssembler::new()
1276 }
1277
1278 pub fn create_zcti_builder(&self, tool_id: u32) -> Result<ZctiBuilder> {
1280 if !self.tools.contains_key(&tool_id) {
1281 return Err(DrivenError::Config(format!("Tool {} not found", tool_id)));
1282 }
1283 Ok(ZctiBuilder::new(tool_id))
1284 }
1285
1286 pub fn invoke_tool_zcti(&self, builder: ZctiBuilder) -> Result<Vec<u8>> {
1288 if !self.is_connected() {
1289 return Err(DrivenError::Config("Not connected".to_string()));
1290 }
1291
1292 let tool_id = builder.tool_id;
1293 if !self.tools.contains_key(&tool_id) {
1294 return Err(DrivenError::Config(format!("Tool {} not found", tool_id)));
1295 }
1296
1297 let (invocation, args_data) = builder.build_with_data();
1298
1299 let mut payload = Vec::with_capacity(ToolInvocation::SIZE + args_data.len());
1301 payload.extend_from_slice(&invocation.as_bytes());
1302 payload.extend_from_slice(&args_data);
1303
1304 let flags = if self.signer.is_some() {
1306 Flags::SIGNED
1307 } else {
1308 0
1309 };
1310 let message = DcpMessage::new(MessageType::Tool, flags, payload);
1311
1312 Ok(message.encode())
1314 }
1315
1316 pub fn parse_tool_invocation<'a>(&self, message: &'a DcpMessage) -> Result<ZctiReader<'a>> {
1318 if message.message_type != MessageType::Tool {
1319 return Err(DrivenError::InvalidBinary(
1320 "Message is not a tool invocation".to_string(),
1321 ));
1322 }
1323
1324 if message.payload.len() < ToolInvocation::SIZE {
1325 return Err(DrivenError::InvalidBinary(
1326 "Payload too small for tool invocation".to_string(),
1327 ));
1328 }
1329
1330 let invocation = ToolInvocation::from_bytes(&message.payload)
1331 .map_err(|e| DrivenError::InvalidBinary(format!("Invalid invocation: {:?}", e)))?;
1332
1333 let args_data = &message.payload[ToolInvocation::SIZE..];
1334 Ok(ZctiReader::new(&invocation, args_data))
1335 }
1336
1337 pub fn create_shared_buffer(&self, capacity: usize) -> SharedArgBuffer {
1339 SharedArgBuffer::new(capacity)
1340 }
1341
1342 pub fn mcp_adapter(&self) -> Option<&McpAdapter> {
1344 self.mcp_adapter.as_ref()
1345 }
1346
1347 pub fn mcp_adapter_mut(&mut self) -> Option<&mut McpAdapter> {
1349 self.mcp_adapter.as_mut()
1350 }
1351
1352 pub fn register_tool_mcp(&mut self, name: &str, tool_id: u32) {
1354 if let Some(ref mut adapter) = self.mcp_adapter {
1355 adapter.register_tool(name, tool_id as u16);
1356 }
1357 }
1358
1359 pub fn handle_mcp_request(&self, json: &str) -> Result<String> {
1363 let adapter = self
1364 .mcp_adapter
1365 .as_ref()
1366 .ok_or_else(|| DrivenError::Config("MCP adapter not available".to_string()))?;
1367
1368 let request = adapter
1370 .parse_request(json)
1371 .map_err(|e| DrivenError::Parse(format!("Invalid JSON-RPC request: {:?}", e)))?;
1372
1373 match request.method.as_str() {
1375 "initialize" => adapter
1376 .handle_initialize(&request)
1377 .map_err(|e| DrivenError::Format(format!("Failed to handle initialize: {:?}", e))),
1378 "tools/list" => adapter
1379 .handle_tools_list(&request)
1380 .map_err(|e| DrivenError::Format(format!("Failed to handle tools/list: {:?}", e))),
1381 "tools/call" => {
1382 self.handle_mcp_tools_call(&request)
1384 }
1385 "resources/list" => {
1386 let result = serde_json::json!({ "resources": [] });
1388 adapter
1389 .format_success_response(request.id, result)
1390 .map_err(|e| DrivenError::Format(format!("Failed to format response: {:?}", e)))
1391 }
1392 "prompts/list" => {
1393 let result = serde_json::json!({ "prompts": [] });
1395 adapter
1396 .format_success_response(request.id, result)
1397 .map_err(|e| DrivenError::Format(format!("Failed to format response: {:?}", e)))
1398 }
1399 _ => {
1400 adapter
1402 .format_error_response(request.id, JsonRpcError::method_not_found())
1403 .map_err(|e| DrivenError::Format(format!("Failed to format error: {:?}", e)))
1404 }
1405 }
1406 }
1407
1408 fn handle_mcp_tools_call(&self, request: &JsonRpcRequest) -> Result<String> {
1410 let adapter = self
1411 .mcp_adapter
1412 .as_ref()
1413 .ok_or_else(|| DrivenError::Config("MCP adapter not available".to_string()))?;
1414
1415 let params = request
1417 .params
1418 .as_ref()
1419 .ok_or_else(|| DrivenError::Parse("Missing params".to_string()))?;
1420
1421 let tool_name = params
1422 .get("name")
1423 .and_then(|v| v.as_str())
1424 .ok_or_else(|| DrivenError::Parse("Missing tool name".to_string()))?;
1425
1426 let arguments = params.get("arguments").cloned();
1427
1428 let tool_id = adapter
1430 .resolve_tool_name(tool_name)
1431 .ok_or_else(|| DrivenError::Config(format!("Unknown tool: {}", tool_name)))?;
1432
1433 if !self.tools.contains_key(&(tool_id as u32)) {
1435 return adapter
1436 .format_error_response(
1437 request.id.clone(),
1438 JsonRpcError::new(-32602, format!("Tool not found: {}", tool_name)),
1439 )
1440 .map_err(|e| DrivenError::Format(format!("Failed to format error: {:?}", e)));
1441 }
1442
1443 let args_bytes = adapter.translate_params(&arguments);
1445
1446 let response_result = serde_json::json!({
1449 "content": [{
1450 "type": "text",
1451 "text": format!("Tool {} called with {} bytes of arguments", tool_name, args_bytes.len())
1452 }]
1453 });
1454
1455 adapter
1456 .format_success_response(request.id.clone(), response_result)
1457 .map_err(|e| DrivenError::Format(format!("Failed to format response: {:?}", e)))
1458 }
1459
1460 pub fn has_mcp_fallback(&self) -> bool {
1462 self.mcp_adapter.is_some()
1463 }
1464
1465 pub fn dcp_to_mcp(&self, message: &DcpMessage) -> Result<String> {
1467 match message.message_type {
1468 MessageType::Response => {
1469 let result: serde_json::Value = serde_json::from_slice(&message.payload)
1471 .unwrap_or_else(|_| {
1472 serde_json::Value::String(
1473 String::from_utf8_lossy(&message.payload).to_string(),
1474 )
1475 });
1476
1477 let response = serde_json::json!({
1478 "jsonrpc": "2.0",
1479 "result": result,
1480 "id": null
1481 });
1482
1483 serde_json::to_string(&response)
1484 .map_err(|e| DrivenError::Format(format!("Failed to serialize: {}", e)))
1485 }
1486 MessageType::Error => {
1487 let (code, msg) = if message.payload.len() >= 4 {
1489 let code_bytes: [u8; 4] = message.payload[0..4].try_into().unwrap();
1490 let code = i32::from_le_bytes(code_bytes);
1491 let msg = String::from_utf8_lossy(&message.payload[4..]).to_string();
1492 (code, msg)
1493 } else {
1494 (-32000, "Unknown error".to_string())
1495 };
1496
1497 let response = serde_json::json!({
1498 "jsonrpc": "2.0",
1499 "error": {
1500 "code": code,
1501 "message": msg
1502 },
1503 "id": null
1504 });
1505
1506 serde_json::to_string(&response)
1507 .map_err(|e| DrivenError::Format(format!("Failed to serialize: {}", e)))
1508 }
1509 _ => Err(DrivenError::Format(format!(
1510 "Cannot convert {:?} message to MCP format",
1511 message.message_type
1512 ))),
1513 }
1514 }
1515
1516 pub fn mcp_to_dcp(&self, json: &str) -> Result<DcpMessage> {
1518 let adapter = self
1519 .mcp_adapter
1520 .as_ref()
1521 .ok_or_else(|| DrivenError::Config("MCP adapter not available".to_string()))?;
1522
1523 let request = adapter
1524 .parse_request(json)
1525 .map_err(|e| DrivenError::Parse(format!("Invalid JSON-RPC: {:?}", e)))?;
1526
1527 let message_type = match request.method.as_str() {
1529 "tools/call" => MessageType::Tool,
1530 "resources/read" => MessageType::Resource,
1531 "prompts/get" => MessageType::Prompt,
1532 _ => MessageType::Tool, };
1534
1535 let payload = serde_json::to_vec(&request)
1537 .map_err(|e| DrivenError::Format(format!("Failed to serialize: {}", e)))?;
1538
1539 Ok(DcpMessage::new(message_type, 0, payload))
1540 }
1541
1542 pub fn invoke_tool_mcp(&self, request: &str) -> Result<String> {
1546 if !self.is_connected() {
1547 return Err(DrivenError::Config("Not connected".to_string()));
1548 }
1549
1550 let _parsed = JsonRpcParser::parse_request(request)
1552 .map_err(|e| DrivenError::Parse(format!("Invalid JSON-RPC request: {:?}", e)))?;
1553
1554 if let Some(ref adapter) = self.mcp_adapter {
1556 let response = adapter
1559 .format_error_response(RequestId::Null, JsonRpcError::method_not_found())
1560 .map_err(|e| DrivenError::Format(format!("Failed to format response: {:?}", e)))?;
1561 Ok(response)
1562 } else {
1563 Err(DrivenError::Config("MCP adapter not available".to_string()))
1564 }
1565 }
1566
1567 pub fn sign_tool_def(&self, tool_id: u32) -> Result<dcp::binary::SignedToolDef> {
1569 let signer = self
1570 .signer
1571 .as_ref()
1572 .ok_or_else(|| DrivenError::Security("Signing not enabled".to_string()))?;
1573
1574 let tool = self
1575 .tools
1576 .get(&tool_id)
1577 .ok_or_else(|| DrivenError::Config(format!("Tool {} not found", tool_id)))?;
1578
1579 Ok(signer.sign_tool_def(tool_id, tool.schema_hash, tool.capabilities))
1580 }
1581
1582 pub fn verify_tool_def(&self, def: &dcp::binary::SignedToolDef) -> Result<()> {
1584 Verifier::verify_tool_def(def)
1585 .map_err(|e| DrivenError::Security(format!("Signature verification failed: {:?}", e)))
1586 }
1587
1588 pub fn sign_invocation(
1593 &self,
1594 tool_id: u32,
1595 nonce: u64,
1596 timestamp: u64,
1597 args: &[u8],
1598 ) -> Result<SignedInvocation> {
1599 let signer = self
1600 .signer
1601 .as_ref()
1602 .ok_or_else(|| DrivenError::Security("Signing not enabled".to_string()))?;
1603
1604 if !self.tools.contains_key(&tool_id) {
1605 return Err(DrivenError::Config(format!("Tool {} not found", tool_id)));
1606 }
1607
1608 Ok(signer.sign_invocation(tool_id, nonce, timestamp, args))
1609 }
1610
1611 pub fn verify_invocation(&self, inv: &SignedInvocation, public_key: &[u8; 32]) -> Result<()> {
1615 Verifier::verify_invocation(inv, public_key)
1616 .map_err(|e| DrivenError::Security(format!("Invocation verification failed: {:?}", e)))
1617 }
1618
1619 pub fn verify_args_hash(&self, inv: &SignedInvocation, args: &[u8]) -> bool {
1621 Verifier::verify_args_hash(inv, args)
1622 }
1623
1624 pub fn generate_signer() -> Signer {
1626 Signer::generate()
1627 }
1628
1629 pub fn parse_capability_manifest(bytes: &[u8]) -> Result<&CapabilityManifest> {
1636 CapabilityManifest::from_bytes(bytes).map_err(|e| {
1637 DrivenError::InvalidBinary(format!("Failed to parse capability manifest: {:?}", e))
1638 })
1639 }
1640
1641 pub fn serialize_capabilities(&self) -> Vec<u8> {
1643 self.capabilities.as_bytes().to_vec()
1644 }
1645
1646 pub fn register_resource(&mut self, resource_id: u16) {
1648 self.capabilities.set_resource(resource_id);
1649 }
1650
1651 pub fn unregister_resource(&mut self, resource_id: u16) {
1653 self.capabilities.clear_resource(resource_id);
1654 }
1655
1656 pub fn has_resource(&self, resource_id: u16) -> bool {
1658 self.capabilities.has_resource(resource_id)
1659 }
1660
1661 pub fn register_prompt(&mut self, prompt_id: u16) {
1663 self.capabilities.set_prompt(prompt_id);
1664 }
1665
1666 pub fn unregister_prompt(&mut self, prompt_id: u16) {
1668 self.capabilities.clear_prompt(prompt_id);
1669 }
1670
1671 pub fn has_prompt(&self, prompt_id: u16) -> bool {
1673 self.capabilities.has_prompt(prompt_id)
1674 }
1675
1676 pub fn set_extension(&mut self, bit: u8) {
1678 self.capabilities.set_extension(bit);
1679 }
1680
1681 pub fn clear_extension(&mut self, bit: u8) {
1683 self.capabilities.clear_extension(bit);
1684 }
1685
1686 pub fn has_extension(&self, bit: u8) -> bool {
1688 self.capabilities.has_extension(bit)
1689 }
1690
1691 pub fn enforce_tool(&self, tool_id: u32) -> Result<()> {
1695 if tool_id >= CapabilityManifest::MAX_TOOLS as u32 {
1696 return Err(DrivenError::Config(format!(
1697 "Tool ID {} exceeds maximum ({})",
1698 tool_id,
1699 CapabilityManifest::MAX_TOOLS
1700 )));
1701 }
1702 if !self.capabilities.has_tool(tool_id as u16) {
1703 return Err(DrivenError::Config(format!(
1704 "Tool {} is not available in capability manifest",
1705 tool_id
1706 )));
1707 }
1708 Ok(())
1709 }
1710
1711 pub fn enforce_resource(&self, resource_id: u16) -> Result<()> {
1715 if !self.capabilities.has_resource(resource_id) {
1716 return Err(DrivenError::Config(format!(
1717 "Resource {} is not available in capability manifest",
1718 resource_id
1719 )));
1720 }
1721 Ok(())
1722 }
1723
1724 pub fn enforce_prompt(&self, prompt_id: u16) -> Result<()> {
1728 if !self.capabilities.has_prompt(prompt_id) {
1729 return Err(DrivenError::Config(format!(
1730 "Prompt {} is not available in capability manifest",
1731 prompt_id
1732 )));
1733 }
1734 Ok(())
1735 }
1736
1737 pub fn negotiate_capabilities(
1742 &mut self,
1743 server_manifest: &CapabilityManifest,
1744 ) -> CapabilityManifest {
1745 let negotiated = self.capabilities.intersect(server_manifest);
1746 self.capabilities = negotiated.clone();
1747 negotiated
1748 }
1749
1750 pub fn capability_stats(&self) -> CapabilityStats {
1752 CapabilityStats {
1753 tool_count: self.capabilities.tool_count(),
1754 resource_count: self.capabilities.resource_count(),
1755 prompt_count: self.capabilities.prompt_count(),
1756 extension_count: self.capabilities.extension_count(),
1757 version: self.capabilities.version,
1758 }
1759 }
1760
1761 pub fn tool_ids(&self) -> impl Iterator<Item = u16> + '_ {
1763 self.capabilities.tool_ids()
1764 }
1765
1766 pub fn resource_ids(&self) -> impl Iterator<Item = u16> + '_ {
1768 self.capabilities.resource_ids()
1769 }
1770
1771 pub fn prompt_ids(&self) -> impl Iterator<Item = u16> + '_ {
1773 self.capabilities.prompt_ids()
1774 }
1775
1776 pub fn intersect_capabilities(&self, other: &CapabilityManifest) -> CapabilityManifest {
1778 self.capabilities.intersect(other)
1779 }
1780
1781 pub fn enable_signing(&mut self, seed: [u8; 32]) {
1783 self.signer = Some(Signer::from_seed(&seed));
1784 self.config.signing_enabled = true;
1785 self.config.signing_seed = Some(seed);
1786 }
1787
1788 pub fn disable_signing(&mut self) {
1790 self.signer = None;
1791 self.config.signing_enabled = false;
1792 self.config.signing_seed = None;
1793 }
1794
1795 pub fn public_key(&self) -> Option<[u8; 32]> {
1797 self.signer.as_ref().map(|s| s.public_key_bytes())
1798 }
1799}
1800
1801pub fn create_envelope(
1803 message_type: MessageType,
1804 flags: u8,
1805 payload_len: u32,
1806) -> BinaryMessageEnvelope {
1807 BinaryMessageEnvelope::new(message_type, flags, payload_len)
1808}
1809
1810pub fn parse_envelope(bytes: &[u8]) -> Result<&BinaryMessageEnvelope> {
1812 BinaryMessageEnvelope::from_bytes(bytes)
1813 .map_err(|e| DrivenError::InvalidBinary(format!("Failed to parse BME: {:?}", e)))
1814}
1815
1816#[cfg(test)]
1817mod tests {
1818 use super::*;
1819 use proptest::prelude::*;
1820
1821 #[test]
1822 fn test_dcp_client_creation() {
1823 let client = DcpClient::with_defaults();
1824 assert_eq!(client.state(), ConnectionState::Disconnected);
1825 assert!(!client.is_connected());
1826 assert!(client.prefer_dcp());
1827 }
1828
1829 #[test]
1830 fn test_dcp_client_config() {
1831 let config = DcpConfig {
1832 enabled: false,
1833 prefer_dcp: false,
1834 endpoint: Some("tcp://localhost:9000".to_string()),
1835 timeout_ms: 10000,
1836 signing_enabled: true,
1837 signing_seed: Some([42u8; 32]),
1838 };
1839
1840 let client = DcpClient::new(config);
1841 assert!(!client.prefer_dcp());
1842 assert!(client.signer.is_some());
1843 }
1844
1845 #[test]
1846 fn test_tool_registration() {
1847 let mut client = DcpClient::with_defaults();
1848
1849 let tool_id = client.register_tool("test_tool", "A test tool", [0xAB; 32], 0x1234);
1850
1851 assert_eq!(tool_id, 1);
1852 assert!(client.get_tool(tool_id).is_some());
1853 assert!(client.capabilities().has_tool(tool_id as u16));
1854
1855 let tool = client.get_tool(tool_id).unwrap();
1856 assert_eq!(tool.name, "test_tool");
1857 assert_eq!(tool.description, "A test tool");
1858 }
1859
1860 #[test]
1861 fn test_signing() {
1862 let mut client = DcpClient::with_defaults();
1863 client.enable_signing([42u8; 32]);
1864
1865 let tool_id = client.register_tool("signed_tool", "A signed tool", [0xCD; 32], 0x5678);
1866
1867 let signed_def = client.sign_tool_def(tool_id).unwrap();
1868 assert!(client.verify_tool_def(&signed_def).is_ok());
1869 }
1870
1871 #[test]
1872 fn test_sign_invocation() {
1873 let mut client = DcpClient::with_defaults();
1874 client.enable_signing([42u8; 32]);
1875
1876 let tool_id = client.register_tool("test_tool", "Test", [0; 32], 0);
1877 let args = b"test arguments";
1878
1879 let signed_inv = client
1880 .sign_invocation(tool_id, 12345, 1234567890, args)
1881 .unwrap();
1882
1883 let public_key = client.public_key().unwrap();
1884 assert!(client.verify_invocation(&signed_inv, &public_key).is_ok());
1885 assert!(client.verify_args_hash(&signed_inv, args));
1886 assert!(!client.verify_args_hash(&signed_inv, b"wrong args"));
1887 }
1888
1889 #[test]
1890 fn test_sign_invocation_without_signing_enabled() {
1891 let mut client = DcpClient::with_defaults();
1892 let tool_id = client.register_tool("test_tool", "Test", [0; 32], 0);
1893
1894 let result = client.sign_invocation(tool_id, 12345, 1234567890, b"args");
1896 assert!(result.is_err());
1897 }
1898
1899 #[test]
1900 fn test_sign_invocation_unknown_tool() {
1901 let mut client = DcpClient::with_defaults();
1902 client.enable_signing([42u8; 32]);
1903
1904 let result = client.sign_invocation(999, 12345, 1234567890, b"args");
1906 assert!(result.is_err());
1907 }
1908
1909 #[test]
1910 fn test_generate_signer() {
1911 let signer1 = DcpClient::generate_signer();
1912 let signer2 = DcpClient::generate_signer();
1913
1914 assert_ne!(signer1.public_key_bytes(), signer2.public_key_bytes());
1916 }
1917
1918 #[test]
1919 fn test_capability_intersection() {
1920 let mut client = DcpClient::with_defaults();
1921 client.register_tool("tool1", "Tool 1", [0; 32], 0);
1922 client.register_tool("tool2", "Tool 2", [0; 32], 0);
1923 client.register_tool("tool3", "Tool 3", [0; 32], 0);
1924
1925 let mut other = CapabilityManifest::new(1);
1926 other.set_tool(2);
1927 other.set_tool(3);
1928 other.set_tool(4);
1929
1930 let intersection = client.intersect_capabilities(&other);
1931 assert!(!intersection.has_tool(1));
1932 assert!(intersection.has_tool(2));
1933 assert!(intersection.has_tool(3));
1934 assert!(!intersection.has_tool(4));
1935 }
1936
1937 #[test]
1938 fn test_envelope_creation() {
1939 let envelope = create_envelope(MessageType::Tool, Flags::STREAMING, 1024);
1940 assert!(envelope.is_streaming());
1941 assert!(!envelope.is_compressed());
1942 assert!(!envelope.is_signed());
1943
1944 let bytes = envelope.as_bytes();
1945 let parsed = parse_envelope(bytes).unwrap();
1946 assert_eq!(parsed.get_message_type(), Some(MessageType::Tool));
1947 }
1948
1949 #[test]
1950 fn test_mcp_fallback() {
1951 let mut client = DcpClient::with_defaults();
1952
1953 let result = client.connect("tcp://localhost:9000");
1955 assert!(result.is_ok());
1956 assert_eq!(client.state(), ConnectionState::ConnectedMcp);
1957 }
1958
1959 #[test]
1960 fn test_protocol_preference() {
1961 let mut client = DcpClient::with_defaults();
1962
1963 assert!(client.prefer_dcp());
1965 assert!(client.is_dcp_enabled());
1966
1967 client.set_prefer_dcp(false);
1969 assert!(!client.prefer_dcp());
1970
1971 client.set_prefer_dcp(true);
1973 assert!(client.prefer_dcp());
1974 }
1975
1976 #[test]
1977 fn test_enable_disable_dcp() {
1978 let mut client = DcpClient::with_defaults();
1979
1980 assert!(client.is_dcp_enabled());
1981
1982 client.disable_dcp();
1983 assert!(!client.is_dcp_enabled());
1984 assert!(!client.prefer_dcp()); client.enable_dcp();
1987 assert!(client.is_dcp_enabled());
1988 }
1989
1990 #[test]
1991 fn test_current_protocol() {
1992 let mut client = DcpClient::with_defaults();
1993
1994 assert!(client.current_protocol().is_none());
1996
1997 client.connect("tcp://localhost:9000").unwrap();
1999 assert_eq!(client.current_protocol(), Some(Protocol::Mcp));
2000 }
2001
2002 #[test]
2003 fn test_select_protocol() {
2004 let mut client = DcpClient::with_defaults();
2005
2006 assert!(client.select_protocol().is_none());
2008
2009 client.connect("tcp://localhost:9000").unwrap();
2011
2012 assert_eq!(client.select_protocol(), Some(Protocol::Mcp));
2014 }
2015
2016 #[test]
2017 fn test_is_mcp_connected() {
2018 let mut client = DcpClient::with_defaults();
2019
2020 assert!(!client.is_mcp_connected());
2021
2022 client.connect("tcp://localhost:9000").unwrap();
2023 assert!(client.is_mcp_connected());
2024 assert!(!client.is_dcp_connected());
2025 }
2026
2027 #[test]
2028 fn test_invoke_tool_auto() {
2029 let mut client = DcpClient::with_defaults();
2030 client.connect("tcp://localhost:9000").unwrap();
2031
2032 let tool_id = client.register_tool("test_tool", "Test", [0; 32], 0);
2033 client.register_tool_mcp("test_tool", tool_id);
2034
2035 let args = b"{}";
2036 let result = client.invoke_tool_auto(tool_id, args);
2037
2038 assert!(result.is_ok());
2040 let result = result.unwrap();
2041 assert!(result.is_mcp());
2042 }
2043
2044 #[test]
2045 fn test_invocation_result() {
2046 let dcp_result = InvocationResult::Dcp(vec![1, 2, 3]);
2047 assert!(dcp_result.is_dcp());
2048 assert!(!dcp_result.is_mcp());
2049 assert_eq!(dcp_result.as_dcp(), Some(&[1, 2, 3][..]));
2050 assert!(dcp_result.as_mcp().is_none());
2051
2052 let mcp_result = InvocationResult::Mcp("{}".to_string());
2053 assert!(!mcp_result.is_dcp());
2054 assert!(mcp_result.is_mcp());
2055 assert!(mcp_result.as_dcp().is_none());
2056 assert_eq!(mcp_result.as_mcp(), Some("{}"));
2057 }
2058
2059 #[test]
2062 fn test_dcp_message_tool() {
2063 let payload = vec![1, 2, 3, 4, 5];
2064 let message = DcpMessage::tool(payload.clone());
2065
2066 assert_eq!(message.message_type, MessageType::Tool);
2067 assert_eq!(message.flags, 0);
2068 assert_eq!(message.payload, payload);
2069 assert!(!message.is_streaming());
2070 assert!(!message.is_compressed());
2071 assert!(!message.is_signed());
2072 }
2073
2074 #[test]
2075 fn test_dcp_message_resource() {
2076 let payload = b"resource://test".to_vec();
2077 let message = DcpMessage::resource(payload.clone());
2078
2079 assert_eq!(message.message_type, MessageType::Resource);
2080 assert_eq!(message.payload, payload);
2081 }
2082
2083 #[test]
2084 fn test_dcp_message_prompt() {
2085 let payload = b"Hello, AI!".to_vec();
2086 let message = DcpMessage::prompt(payload.clone());
2087
2088 assert_eq!(message.message_type, MessageType::Prompt);
2089 assert_eq!(message.payload, payload);
2090 }
2091
2092 #[test]
2093 fn test_dcp_message_response() {
2094 let payload = b"Response data".to_vec();
2095 let message = DcpMessage::response(payload.clone());
2096
2097 assert_eq!(message.message_type, MessageType::Response);
2098 assert_eq!(message.payload, payload);
2099 }
2100
2101 #[test]
2102 fn test_dcp_message_error() {
2103 let payload = b"Error occurred".to_vec();
2104 let message = DcpMessage::error(payload.clone());
2105
2106 assert_eq!(message.message_type, MessageType::Error);
2107 assert_eq!(message.payload, payload);
2108 }
2109
2110 #[test]
2111 fn test_dcp_message_stream() {
2112 let payload = b"Streaming data".to_vec();
2113 let message = DcpMessage::stream(payload.clone());
2114
2115 assert_eq!(message.message_type, MessageType::Stream);
2116 assert!(message.is_streaming());
2117 assert_eq!(message.payload, payload);
2118 }
2119
2120 #[test]
2121 fn test_dcp_message_flags() {
2122 let message = DcpMessage::tool(vec![])
2123 .with_streaming()
2124 .with_compressed()
2125 .with_signed();
2126
2127 assert!(message.is_streaming());
2128 assert!(message.is_compressed());
2129 assert!(message.is_signed());
2130 }
2131
2132 #[test]
2133 fn test_dcp_message_encode_decode_roundtrip() {
2134 let original = DcpMessage::new(
2135 MessageType::Tool,
2136 Flags::STREAMING | Flags::SIGNED,
2137 vec![0xDE, 0xAD, 0xBE, 0xEF],
2138 );
2139
2140 let encoded = original.encode();
2141 let decoded = DcpMessage::decode(&encoded).unwrap();
2142
2143 assert_eq!(decoded.message_type, original.message_type);
2144 assert_eq!(decoded.flags, original.flags);
2145 assert_eq!(decoded.payload, original.payload);
2146 }
2147
2148 #[test]
2149 fn test_dcp_message_decode_insufficient_data() {
2150 let bytes = vec![0u8; 4]; let result = DcpMessage::decode(&bytes);
2152 assert!(result.is_err());
2153 }
2154
2155 #[test]
2156 fn test_dcp_message_decode_invalid_magic() {
2157 let mut bytes = vec![0u8; 16];
2158 bytes[0] = 0xFF;
2159 bytes[1] = 0xFF;
2160 let result = DcpMessage::decode(&bytes);
2161 assert!(result.is_err());
2162 }
2163
2164 #[test]
2165 fn test_dcp_message_all_types_roundtrip() {
2166 let test_cases = vec![
2167 DcpMessage::tool(vec![1, 2, 3]),
2168 DcpMessage::resource(b"uri://test".to_vec()),
2169 DcpMessage::prompt(b"test prompt".to_vec()),
2170 DcpMessage::response(b"test response".to_vec()),
2171 DcpMessage::error(b"test error".to_vec()),
2172 DcpMessage::stream(b"streaming".to_vec()),
2173 ];
2174
2175 for original in test_cases {
2176 let encoded = original.encode();
2177 let decoded = DcpMessage::decode(&encoded).unwrap();
2178 assert_eq!(decoded.message_type, original.message_type);
2179 assert_eq!(decoded.payload, original.payload);
2180 }
2181 }
2182
2183 #[test]
2186 fn test_stream_assembler_single_chunk() {
2187 let mut assembler = StreamAssembler::new();
2188
2189 let chunk = StreamChunk::new(0, ChunkFlags::FIRST | ChunkFlags::LAST, 5);
2190 let payload = vec![1, 2, 3, 4, 5];
2191
2192 assembler.add_chunk(chunk.as_bytes(), &payload).unwrap();
2193
2194 assert!(assembler.is_complete());
2195 assert!(!assembler.has_error());
2196 assert_eq!(assembler.data(), Some(payload.as_slice()));
2197 assert_eq!(assembler.chunk_count(), 1);
2198 }
2199
2200 #[test]
2201 fn test_stream_assembler_multiple_chunks() {
2202 let mut assembler = StreamAssembler::new();
2203
2204 let chunk1 = StreamChunk::first(0, 3);
2206 assembler.add_chunk(chunk1.as_bytes(), &[1, 2, 3]).unwrap();
2207 assert!(!assembler.is_complete());
2208
2209 let chunk2 = StreamChunk::continuation(1, 3);
2211 assembler.add_chunk(chunk2.as_bytes(), &[4, 5, 6]).unwrap();
2212 assert!(!assembler.is_complete());
2213
2214 let chunk3 = StreamChunk::last(2, 2);
2216 assembler.add_chunk(chunk3.as_bytes(), &[7, 8]).unwrap();
2217 assert!(assembler.is_complete());
2218
2219 assert_eq!(assembler.data(), Some(&[1, 2, 3, 4, 5, 6, 7, 8][..]));
2220 assert_eq!(assembler.chunk_count(), 3);
2221 }
2222
2223 #[test]
2224 fn test_stream_assembler_out_of_order() {
2225 let mut assembler = StreamAssembler::new();
2226
2227 let chunk1 = StreamChunk::first(0, 3);
2228 assembler.add_chunk(chunk1.as_bytes(), &[1, 2, 3]).unwrap();
2229
2230 let chunk_wrong = StreamChunk::continuation(5, 3);
2232 let result = assembler.add_chunk(chunk_wrong.as_bytes(), &[4, 5, 6]);
2233 assert!(result.is_err());
2234 }
2235
2236 #[test]
2237 fn test_stream_assembler_error_chunk() {
2238 let mut assembler = StreamAssembler::new();
2239
2240 let chunk = StreamChunk::error(0, 0);
2241 let result = assembler.add_chunk(chunk.as_bytes(), &[]);
2242
2243 assert!(result.is_err());
2244 assert!(assembler.has_error());
2245 assert!(assembler.is_complete());
2246 assert!(assembler.data().is_none());
2247 }
2248
2249 #[test]
2250 fn test_stream_assembler_take_data() {
2251 let mut assembler = StreamAssembler::new();
2252
2253 let chunk = StreamChunk::new(0, ChunkFlags::FIRST | ChunkFlags::LAST, 3);
2254 assembler.add_chunk(chunk.as_bytes(), &[1, 2, 3]).unwrap();
2255
2256 let data = assembler.take_data();
2257 assert_eq!(data, Some(vec![1, 2, 3]));
2258 }
2259
2260 #[test]
2263 fn test_stream_builder_single_chunk() {
2264 let mut builder = StreamBuilder::new(100);
2265 let data = vec![1, 2, 3, 4, 5];
2266
2267 let chunks = builder.build_chunks(&data);
2268
2269 assert_eq!(chunks.len(), 1);
2270 let (chunk, payload) = &chunks[0];
2271 assert!(chunk.is_first());
2272 assert!(chunk.is_last());
2273 assert_eq!(payload, &data);
2274 }
2275
2276 #[test]
2277 fn test_stream_builder_multiple_chunks() {
2278 let mut builder = StreamBuilder::new(3);
2279 let data = vec![1, 2, 3, 4, 5, 6, 7, 8];
2280
2281 let chunks = builder.build_chunks(&data);
2282
2283 assert_eq!(chunks.len(), 3);
2284
2285 assert!(chunks[0].0.is_first());
2287 assert!(!chunks[0].0.is_last());
2288 assert_eq!(chunks[0].1, vec![1, 2, 3]);
2289
2290 assert!(chunks[1].0.is_continuation());
2292 assert!(!chunks[1].0.is_last());
2293 assert_eq!(chunks[1].1, vec![4, 5, 6]);
2294
2295 assert!(!chunks[2].0.is_first());
2297 assert!(chunks[2].0.is_last());
2298 assert_eq!(chunks[2].1, vec![7, 8]);
2299 }
2300
2301 #[test]
2302 fn test_stream_builder_reset() {
2303 let mut builder = StreamBuilder::new(10);
2304
2305 let _ = builder.build_chunks(&[1, 2, 3]);
2306 builder.reset();
2307
2308 let chunks = builder.build_chunks(&[4, 5, 6]);
2309 let sequence = chunks[0].0.sequence;
2311 assert_eq!(sequence, 0);
2312 }
2313
2314 #[test]
2315 fn test_stream_builder_assembler_roundtrip() {
2316 let mut builder = StreamBuilder::new(5);
2317 let original_data = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12];
2318
2319 let chunks = builder.build_chunks(&original_data);
2320
2321 let mut assembler = StreamAssembler::new();
2322 for (chunk, payload) in chunks {
2323 assembler.add_chunk(chunk.as_bytes(), &payload).unwrap();
2324 }
2325
2326 assert!(assembler.is_complete());
2327 assert_eq!(assembler.take_data(), Some(original_data));
2328 }
2329
2330 #[test]
2333 fn test_client_create_tool_message() {
2334 let mut client = DcpClient::with_defaults();
2335 client.connect("tcp://localhost:9000").unwrap();
2336
2337 let tool_id = client.register_tool("test", "Test tool", [0; 32], 0);
2338 let args = vec![1, 2, 3, 4];
2339
2340 let message = client.create_tool_message(tool_id, &args).unwrap();
2341
2342 assert_eq!(message.message_type, MessageType::Tool);
2343 assert!(message.payload.len() >= ToolInvocation::SIZE + args.len());
2344 }
2345
2346 #[test]
2347 fn test_client_create_resource_message() {
2348 let client = DcpClient::with_defaults();
2349 let message = client.create_resource_message("file:///test.txt");
2350
2351 assert_eq!(message.message_type, MessageType::Resource);
2352 assert_eq!(message.payload, b"file:///test.txt");
2353 }
2354
2355 #[test]
2356 fn test_client_create_prompt_message() {
2357 let client = DcpClient::with_defaults();
2358 let message = client.create_prompt_message("Hello, AI!");
2359
2360 assert_eq!(message.message_type, MessageType::Prompt);
2361 assert_eq!(message.payload, b"Hello, AI!");
2362 }
2363
2364 #[test]
2365 fn test_client_create_error_message() {
2366 let client = DcpClient::with_defaults();
2367 let message = client.create_error_message(404, "Not found");
2368
2369 assert_eq!(message.message_type, MessageType::Error);
2370 assert_eq!(&message.payload[0..4], &404u32.to_le_bytes());
2372 assert_eq!(&message.payload[4..], b"Not found");
2373 }
2374
2375 #[test]
2376 fn test_client_encode_decode_message() {
2377 let client = DcpClient::with_defaults();
2378
2379 let payload = b"test payload".to_vec();
2380 let encoded = client.encode_message(MessageType::Response, &payload);
2381 let decoded = client.decode_message(&encoded).unwrap();
2382
2383 assert_eq!(decoded.message_type, MessageType::Response);
2384 assert_eq!(decoded.payload, payload);
2385 }
2386
2387 #[test]
2388 fn test_client_stream_builder_assembler() {
2389 let client = DcpClient::with_defaults();
2390
2391 let mut builder = client.create_stream_builder(10);
2392 let assembler = client.create_stream_assembler();
2393
2394 let chunks = builder.build_chunks(&[1, 2, 3]);
2396 assert_eq!(chunks.len(), 1);
2397 assert!(!assembler.is_complete());
2398 }
2399
2400 #[test]
2403 fn test_zcti_builder_basic() {
2404 let builder = ZctiBuilder::new(42)
2405 .add_bool(true)
2406 .add_i32(123)
2407 .add_string("hello");
2408
2409 assert_eq!(builder.arg_count(), 3);
2410
2411 let (invocation, data) = builder.build_with_data();
2412 assert_eq!(invocation.tool_id, 42);
2413 assert_eq!(invocation.arg_count(), 3);
2414 assert!(!data.is_empty());
2415 }
2416
2417 #[test]
2418 fn test_zcti_builder_all_types() {
2419 let builder = ZctiBuilder::new(1)
2420 .add_null()
2421 .add_bool(false)
2422 .add_i32(-42)
2423 .add_i64(i64::MAX)
2424 .add_f64(3.14159)
2425 .add_string("test")
2426 .add_bytes(&[0xDE, 0xAD, 0xBE, 0xEF]);
2427
2428 assert_eq!(builder.arg_count(), 7);
2429 }
2430
2431 #[test]
2432 fn test_zcti_reader_basic() {
2433 let builder = ZctiBuilder::new(42)
2434 .add_bool(true)
2435 .add_i32(123)
2436 .add_string("hello");
2437
2438 let (invocation, data) = builder.build_with_data();
2439 let mut reader = ZctiReader::new(&invocation, &data);
2440
2441 assert_eq!(reader.tool_id(), 42);
2442 assert_eq!(reader.arg_count(), 3);
2443 assert!(reader.has_more());
2444
2445 assert_eq!(reader.read_bool().unwrap(), true);
2446 assert_eq!(reader.read_i32().unwrap(), 123);
2447 assert_eq!(reader.read_string().unwrap(), "hello");
2448 assert!(!reader.has_more());
2449 }
2450
2451 #[test]
2452 fn test_zcti_reader_all_types() {
2453 let builder = ZctiBuilder::new(1)
2454 .add_bool(false)
2455 .add_i32(-42)
2456 .add_i64(i64::MAX)
2457 .add_f64(3.14159)
2458 .add_string("test")
2459 .add_bytes(&[0xDE, 0xAD, 0xBE, 0xEF]);
2460
2461 let (invocation, data) = builder.build_with_data();
2462 let mut reader = ZctiReader::new(&invocation, &data);
2463
2464 assert_eq!(reader.read_bool().unwrap(), false);
2465 assert_eq!(reader.read_i32().unwrap(), -42);
2466 assert_eq!(reader.read_i64().unwrap(), i64::MAX);
2467 assert!((reader.read_f64().unwrap() - 3.14159).abs() < 0.00001);
2468 assert_eq!(reader.read_string().unwrap(), "test");
2469 assert_eq!(reader.read_bytes().unwrap(), &[0xDE, 0xAD, 0xBE, 0xEF]);
2470 }
2471
2472 #[test]
2473 fn test_zcti_reader_skip() {
2474 let builder = ZctiBuilder::new(1)
2475 .add_bool(true)
2476 .add_i32(123)
2477 .add_string("skip me")
2478 .add_i64(456);
2479
2480 let (invocation, data) = builder.build_with_data();
2481 let mut reader = ZctiReader::new(&invocation, &data);
2482
2483 reader.skip().unwrap(); reader.skip().unwrap(); reader.skip().unwrap(); assert_eq!(reader.read_i64().unwrap(), 456);
2487 }
2488
2489 #[test]
2490 fn test_zcti_type_mismatch() {
2491 let builder = ZctiBuilder::new(1).add_bool(true);
2492 let (invocation, data) = builder.build_with_data();
2493 let mut reader = ZctiReader::new(&invocation, &data);
2494
2495 let result = reader.read_i32();
2497 assert!(result.is_err());
2498 }
2499
2500 #[test]
2501 fn test_shared_arg_buffer() {
2502 let mut buffer = SharedArgBuffer::new(1024);
2503
2504 let data1 = b"hello";
2505 let offset1 = buffer.write(data1).unwrap();
2506 assert_eq!(offset1, 0);
2507
2508 let data2 = b"world";
2509 let offset2 = buffer.write(data2).unwrap();
2510 assert_eq!(offset2, 5);
2511
2512 assert_eq!(buffer.read(offset1, 5).unwrap(), b"hello");
2513 assert_eq!(buffer.read(offset2, 5).unwrap(), b"world");
2514 }
2515
2516 #[test]
2517 fn test_shared_arg_buffer_overflow() {
2518 let mut buffer = SharedArgBuffer::new(10);
2519
2520 let result = buffer.write(&[0u8; 20]);
2521 assert!(result.is_err());
2522 }
2523
2524 #[test]
2525 fn test_shared_arg_buffer_reset() {
2526 let mut buffer = SharedArgBuffer::new(100);
2527
2528 buffer.write(b"test").unwrap();
2529 assert_eq!(buffer.offset(), 4);
2530
2531 buffer.reset();
2532 assert_eq!(buffer.offset(), 0);
2533 }
2534
2535 #[test]
2536 fn test_client_zcti_integration() {
2537 let mut client = DcpClient::with_defaults();
2538 client.connect("tcp://localhost:9000").unwrap();
2539
2540 let tool_id = client.register_tool("test_tool", "Test", [0; 32], 0);
2541
2542 let builder = client
2543 .create_zcti_builder(tool_id)
2544 .unwrap()
2545 .add_string("arg1")
2546 .add_i32(42);
2547
2548 let encoded = client.invoke_tool_zcti(builder).unwrap();
2549
2550 let message = client.decode_message(&encoded).unwrap();
2552 assert_eq!(message.message_type, MessageType::Tool);
2553
2554 let mut reader = client.parse_tool_invocation(&message).unwrap();
2555 assert_eq!(reader.tool_id(), tool_id);
2556 assert_eq!(reader.read_string().unwrap(), "arg1");
2557 assert_eq!(reader.read_i32().unwrap(), 42);
2558 }
2559
2560 #[test]
2563 fn test_mcp_adapter_available() {
2564 let mut client = DcpClient::with_defaults();
2565 client.connect("tcp://localhost:9000").unwrap();
2566
2567 assert!(client.has_mcp_fallback());
2568 assert!(client.mcp_adapter().is_some());
2569 }
2570
2571 #[test]
2572 fn test_mcp_register_tool() {
2573 let mut client = DcpClient::with_defaults();
2574 client.connect("tcp://localhost:9000").unwrap();
2575
2576 let tool_id = client.register_tool("read_file", "Read a file", [0; 32], 0);
2577 client.register_tool_mcp("read_file", tool_id);
2578
2579 let adapter = client.mcp_adapter().unwrap();
2580 assert_eq!(adapter.resolve_tool_name("read_file"), Some(tool_id as u16));
2581 }
2582
2583 #[test]
2584 fn test_mcp_handle_initialize() {
2585 let mut client = DcpClient::with_defaults();
2586 client.connect("tcp://localhost:9000").unwrap();
2587
2588 let request = r#"{"jsonrpc":"2.0","method":"initialize","id":1}"#;
2589 let response = client.handle_mcp_request(request).unwrap();
2590
2591 assert!(response.contains("protocolVersion"));
2592 assert!(response.contains("capabilities"));
2593 }
2594
2595 #[test]
2596 fn test_mcp_handle_tools_list() {
2597 let mut client = DcpClient::with_defaults();
2598 client.connect("tcp://localhost:9000").unwrap();
2599
2600 let tool_id = client.register_tool("test_tool", "Test tool", [0; 32], 0);
2601 client.register_tool_mcp("test_tool", tool_id);
2602
2603 let request = r#"{"jsonrpc":"2.0","method":"tools/list","id":2}"#;
2604 let response = client.handle_mcp_request(request).unwrap();
2605
2606 assert!(response.contains("tools"));
2607 assert!(response.contains("test_tool"));
2608 }
2609
2610 #[test]
2611 fn test_mcp_handle_tools_call() {
2612 let mut client = DcpClient::with_defaults();
2613 client.connect("tcp://localhost:9000").unwrap();
2614
2615 let tool_id = client.register_tool("echo", "Echo tool", [0; 32], 0);
2616 client.register_tool_mcp("echo", tool_id);
2617
2618 let request = r#"{"jsonrpc":"2.0","method":"tools/call","params":{"name":"echo","arguments":{"text":"hello"}},"id":3}"#;
2619 let response = client.handle_mcp_request(request).unwrap();
2620
2621 assert!(response.contains("content"));
2622 assert!(response.contains("echo"));
2623 }
2624
2625 #[test]
2626 fn test_mcp_handle_unknown_method() {
2627 let mut client = DcpClient::with_defaults();
2628 client.connect("tcp://localhost:9000").unwrap();
2629
2630 let request = r#"{"jsonrpc":"2.0","method":"unknown/method","id":4}"#;
2631 let response = client.handle_mcp_request(request).unwrap();
2632
2633 assert!(response.contains("error"));
2634 assert!(response.contains("-32601")); }
2636
2637 #[test]
2638 fn test_mcp_handle_resources_list() {
2639 let mut client = DcpClient::with_defaults();
2640 client.connect("tcp://localhost:9000").unwrap();
2641
2642 let request = r#"{"jsonrpc":"2.0","method":"resources/list","id":5}"#;
2643 let response = client.handle_mcp_request(request).unwrap();
2644
2645 assert!(response.contains("resources"));
2646 assert!(response.contains("[]")); }
2648
2649 #[test]
2650 fn test_mcp_handle_prompts_list() {
2651 let mut client = DcpClient::with_defaults();
2652 client.connect("tcp://localhost:9000").unwrap();
2653
2654 let request = r#"{"jsonrpc":"2.0","method":"prompts/list","id":6}"#;
2655 let response = client.handle_mcp_request(request).unwrap();
2656
2657 assert!(response.contains("prompts"));
2658 assert!(response.contains("[]")); }
2660
2661 #[test]
2662 fn test_dcp_to_mcp_response() {
2663 let client = DcpClient::with_defaults();
2664
2665 let payload = serde_json::to_vec(&serde_json::json!({"result": "success"})).unwrap();
2666 let message = DcpMessage::response(payload);
2667
2668 let mcp_json = client.dcp_to_mcp(&message).unwrap();
2669 assert!(mcp_json.contains("jsonrpc"));
2670 assert!(mcp_json.contains("result"));
2671 }
2672
2673 #[test]
2674 fn test_dcp_to_mcp_error() {
2675 let client = DcpClient::with_defaults();
2676
2677 let message = client.create_error_message(404, "Not found");
2678 let mcp_json = client.dcp_to_mcp(&message).unwrap();
2679
2680 assert!(mcp_json.contains("error"));
2681 assert!(mcp_json.contains("404"));
2682 assert!(mcp_json.contains("Not found"));
2683 }
2684
2685 #[test]
2686 fn test_mcp_to_dcp() {
2687 let mut client = DcpClient::with_defaults();
2688 client.connect("tcp://localhost:9000").unwrap();
2689
2690 let mcp_request =
2691 r#"{"jsonrpc":"2.0","method":"tools/call","params":{"name":"test"},"id":1}"#;
2692 let dcp_message = client.mcp_to_dcp(mcp_request).unwrap();
2693
2694 assert_eq!(dcp_message.message_type, MessageType::Tool);
2695 assert!(!dcp_message.payload.is_empty());
2696 }
2697
2698 fn arb_message_type() -> impl Strategy<Value = MessageType> {
2702 prop_oneof![
2703 Just(MessageType::Tool),
2704 Just(MessageType::Resource),
2705 Just(MessageType::Prompt),
2706 Just(MessageType::Response),
2707 Just(MessageType::Error),
2708 Just(MessageType::Stream),
2709 ]
2710 }
2711
2712 fn arb_flags() -> impl Strategy<Value = u8> {
2714 prop_oneof![
2715 Just(0u8),
2716 Just(Flags::STREAMING),
2717 Just(Flags::COMPRESSED),
2718 Just(Flags::SIGNED),
2719 Just(Flags::STREAMING | Flags::COMPRESSED),
2720 Just(Flags::STREAMING | Flags::SIGNED),
2721 Just(Flags::COMPRESSED | Flags::SIGNED),
2722 Just(Flags::STREAMING | Flags::COMPRESSED | Flags::SIGNED),
2723 ]
2724 }
2725
2726 fn arb_payload() -> impl Strategy<Value = Vec<u8>> {
2728 prop::collection::vec(any::<u8>(), 0..1024)
2729 }
2730
2731 fn arb_dcp_message() -> impl Strategy<Value = DcpMessage> {
2733 (arb_message_type(), arb_flags(), arb_payload()).prop_map(
2734 |(message_type, flags, payload)| DcpMessage::new(message_type, flags, payload),
2735 )
2736 }
2737
2738 fn arb_chunk_size() -> impl Strategy<Value = usize> {
2740 1usize..256
2741 }
2742
2743 fn arb_stream_data() -> impl Strategy<Value = Vec<u8>> {
2745 prop::collection::vec(any::<u8>(), 1..2048)
2746 }
2747
2748 proptest! {
2749 #![proptest_config(ProptestConfig::with_cases(100))]
2750
2751 #[test]
2756 fn prop_bme_roundtrip(message in arb_dcp_message()) {
2757 let encoded = message.encode();
2758 let decoded = DcpMessage::decode(&encoded).expect("Decoding should succeed");
2759
2760 prop_assert_eq!(decoded.message_type, message.message_type);
2761 prop_assert_eq!(decoded.flags, message.flags);
2762 prop_assert_eq!(decoded.payload, message.payload);
2763 }
2764
2765 #[test]
2769 fn prop_bme_envelope_header(
2770 message_type in arb_message_type(),
2771 flags in arb_flags(),
2772 payload_len in 0u32..65536
2773 ) {
2774 let envelope = BinaryMessageEnvelope::new(message_type, flags, payload_len);
2775 let bytes = envelope.as_bytes();
2776 let parsed = BinaryMessageEnvelope::from_bytes(bytes).expect("Parsing should succeed");
2777
2778 let parsed_type = parsed.message_type;
2780 let parsed_flags = parsed.flags;
2781 let parsed_len = parsed.payload_len;
2782
2783 prop_assert_eq!(parsed_type, message_type as u8);
2784 prop_assert_eq!(parsed_flags, flags);
2785 prop_assert_eq!(parsed_len, payload_len);
2786 }
2787
2788 #[test]
2792 fn prop_stream_roundtrip(
2793 data in arb_stream_data(),
2794 chunk_size in arb_chunk_size()
2795 ) {
2796 let mut builder = StreamBuilder::new(chunk_size);
2797 let chunks = builder.build_chunks(&data);
2798
2799 let mut assembler = StreamAssembler::new();
2800 for (chunk, payload) in chunks {
2801 assembler.add_chunk(chunk.as_bytes(), &payload).expect("Adding chunk should succeed");
2802 }
2803
2804 prop_assert!(assembler.is_complete());
2805 prop_assert!(!assembler.has_error());
2806 prop_assert_eq!(assembler.take_data(), Some(data));
2807 }
2808
2809 #[test]
2813 fn prop_all_message_types_roundtrip(
2814 message_type in arb_message_type(),
2815 payload in arb_payload()
2816 ) {
2817 let message = DcpMessage::new(message_type, 0, payload.clone());
2818 let encoded = message.encode();
2819 let decoded = DcpMessage::decode(&encoded).expect("Decoding should succeed");
2820
2821 prop_assert_eq!(decoded.message_type, message_type);
2822 prop_assert_eq!(decoded.payload, payload);
2823 }
2824
2825 #[test]
2829 fn prop_flags_preserved(flags in arb_flags()) {
2830 let message = DcpMessage::new(MessageType::Tool, flags, vec![1, 2, 3]);
2831 let encoded = message.encode();
2832 let decoded = DcpMessage::decode(&encoded).expect("Decoding should succeed");
2833
2834 prop_assert_eq!(decoded.flags, flags);
2835 prop_assert_eq!(decoded.is_streaming(), flags & Flags::STREAMING != 0);
2836 prop_assert_eq!(decoded.is_compressed(), flags & Flags::COMPRESSED != 0);
2837 prop_assert_eq!(decoded.is_signed(), flags & Flags::SIGNED != 0);
2838 }
2839
2840 #[test]
2845 fn prop_mcp_initialize_response_valid(id in 1i64..1000) {
2846 let request = format!(r#"{{"jsonrpc":"2.0","method":"initialize","id":{}}}"#, id);
2847
2848 let mut client = DcpClient::with_defaults();
2849 client.connect("tcp://localhost:9000").expect("Connection should succeed");
2850
2851 let response = client.handle_mcp_request(&request).expect("Should handle initialize");
2852
2853 let parsed: serde_json::Value = serde_json::from_str(&response).expect("Should be valid JSON");
2855 prop_assert_eq!(parsed["jsonrpc"].as_str(), Some("2.0"));
2856 prop_assert!(parsed.get("result").is_some() || parsed.get("error").is_some());
2857
2858 if let Some(result) = parsed.get("result") {
2860 prop_assert!(result.get("capabilities").is_some());
2861 prop_assert!(result.get("protocolVersion").is_some());
2862 }
2863 }
2864
2865 #[test]
2868 fn prop_mcp_tools_list_contains_registered(
2869 tool_names in prop::collection::vec("[a-z_]+", 1..5)
2870 ) {
2871 let mut client = DcpClient::with_defaults();
2872 client.connect("tcp://localhost:9000").expect("Connection should succeed");
2873
2874 for name in &tool_names {
2876 let tool_id = client.register_tool(name, "Test tool", [0; 32], 0);
2877 client.register_tool_mcp(name, tool_id);
2878 }
2879
2880 let request = r#"{"jsonrpc":"2.0","method":"tools/list","id":1}"#;
2881 let response = client.handle_mcp_request(request).expect("Should handle tools/list");
2882
2883 let parsed: serde_json::Value = serde_json::from_str(&response).expect("Should be valid JSON");
2884 let tools = parsed["result"]["tools"].as_array().expect("Should have tools array");
2885
2886 for name in &tool_names {
2888 let found = tools.iter().any(|t| t["name"].as_str() == Some(name.as_str()));
2889 prop_assert!(found, "Tool {} should be in response", name);
2890 }
2891 }
2892
2893 #[test]
2898 fn prop_capability_intersection(
2899 tools1 in prop::collection::vec(0u16..1000, 0..20),
2900 tools2 in prop::collection::vec(0u16..1000, 0..20),
2901 resources1 in prop::collection::vec(0u16..500, 0..10),
2902 resources2 in prop::collection::vec(0u16..500, 0..10),
2903 prompts1 in prop::collection::vec(0u16..200, 0..5),
2904 prompts2 in prop::collection::vec(0u16..200, 0..5),
2905 ) {
2906 let mut m1 = CapabilityManifest::new(1);
2907 let mut m2 = CapabilityManifest::new(1);
2908
2909 for &t in &tools1 { m1.set_tool(t); }
2911 for &r in &resources1 { m1.set_resource(r); }
2912 for &p in &prompts1 { m1.set_prompt(p); }
2913
2914 for &t in &tools2 { m2.set_tool(t); }
2916 for &r in &resources2 { m2.set_resource(r); }
2917 for &p in &prompts2 { m2.set_prompt(p); }
2918
2919 let intersection = m1.intersect(&m2);
2921
2922 for t in 0u16..1000 {
2924 let in_both = m1.has_tool(t) && m2.has_tool(t);
2925 prop_assert_eq!(
2926 intersection.has_tool(t), in_both,
2927 "Tool {} should be in intersection iff in both manifests", t
2928 );
2929 }
2930
2931 for r in 0u16..500 {
2932 let in_both = m1.has_resource(r) && m2.has_resource(r);
2933 prop_assert_eq!(
2934 intersection.has_resource(r), in_both,
2935 "Resource {} should be in intersection iff in both manifests", r
2936 );
2937 }
2938
2939 for p in 0u16..200 {
2940 let in_both = m1.has_prompt(p) && m2.has_prompt(p);
2941 prop_assert_eq!(
2942 intersection.has_prompt(p), in_both,
2943 "Prompt {} should be in intersection iff in both manifests", p
2944 );
2945 }
2946 }
2947
2948 #[test]
2952 fn prop_capability_manifest_roundtrip(
2953 tools in prop::collection::vec(0u16..8000, 0..50),
2954 resources in prop::collection::vec(0u16..1000, 0..20),
2955 prompts in prop::collection::vec(0u16..500, 0..10),
2956 extensions in 0u64..u64::MAX,
2957 ) {
2958 let mut manifest = CapabilityManifest::new(1);
2959
2960 for &t in &tools { manifest.set_tool(t); }
2961 for &r in &resources { manifest.set_resource(r); }
2962 for &p in &prompts { manifest.set_prompt(p); }
2963 manifest.extensions = extensions;
2964
2965 let bytes = manifest.as_bytes();
2967 let parsed = CapabilityManifest::from_bytes(bytes).expect("Should parse");
2968
2969 for &t in &tools {
2971 prop_assert!(parsed.has_tool(t), "Tool {} should be preserved", t);
2972 }
2973 for &r in &resources {
2974 prop_assert!(parsed.has_resource(r), "Resource {} should be preserved", r);
2975 }
2976 for &p in &prompts {
2977 prop_assert!(parsed.has_prompt(p), "Prompt {} should be preserved", p);
2978 }
2979 prop_assert_eq!(parsed.extensions, extensions);
2980 }
2981
2982 #[test]
2987 fn prop_signature_verification_tool_def(
2988 seed in prop::array::uniform32(any::<u8>()),
2989 tool_id in 1u32..1000,
2990 schema_hash in prop::array::uniform32(any::<u8>()),
2991 capabilities in any::<u64>(),
2992 ) {
2993 let signer = Signer::from_seed(&seed);
2994 let signed_def = signer.sign_tool_def(tool_id, schema_hash, capabilities);
2995
2996 prop_assert!(Verifier::verify_tool_def(&signed_def).is_ok());
2998
2999 let mut tampered = signed_def.clone();
3001 tampered.tool_id = tool_id.wrapping_add(1);
3002 prop_assert!(Verifier::verify_tool_def(&tampered).is_err());
3003
3004 let mut tampered = signed_def.clone();
3005 tampered.schema_hash[0] ^= 0xFF;
3006 prop_assert!(Verifier::verify_tool_def(&tampered).is_err());
3007
3008 let mut tampered = signed_def.clone();
3009 tampered.capabilities ^= 0xFFFFFFFF;
3010 prop_assert!(Verifier::verify_tool_def(&tampered).is_err());
3011 }
3012
3013 #[test]
3017 fn prop_signature_verification_invocation(
3018 seed in prop::array::uniform32(any::<u8>()),
3019 wrong_seed in prop::array::uniform32(any::<u8>()),
3020 tool_id in 1u32..1000,
3021 nonce in any::<u64>(),
3022 timestamp in any::<u64>(),
3023 args in prop::collection::vec(any::<u8>(), 0..256),
3024 ) {
3025 let signer = Signer::from_seed(&seed);
3026 let public_key = signer.public_key_bytes();
3027 let signed_inv = signer.sign_invocation(tool_id, nonce, timestamp, &args);
3028
3029 prop_assert!(Verifier::verify_invocation(&signed_inv, &public_key).is_ok());
3031
3032 prop_assert!(Verifier::verify_args_hash(&signed_inv, &args));
3034
3035 if seed != wrong_seed {
3037 let wrong_signer = Signer::from_seed(&wrong_seed);
3038 let wrong_key = wrong_signer.public_key_bytes();
3039 prop_assert!(Verifier::verify_invocation(&signed_inv, &wrong_key).is_err());
3040 }
3041
3042 let mut tampered = signed_inv.clone();
3044 tampered.tool_id = tool_id.wrapping_add(1);
3045 prop_assert!(Verifier::verify_invocation(&tampered, &public_key).is_err());
3046
3047 let mut tampered = signed_inv.clone();
3048 tampered.nonce ^= 0xFFFFFFFF;
3049 prop_assert!(Verifier::verify_invocation(&tampered, &public_key).is_err());
3050 }
3051 }
3052
3053 #[test]
3056 fn test_capability_manifest_parsing() {
3057 let mut manifest = CapabilityManifest::new(1);
3058 manifest.set_tool(42);
3059 manifest.set_resource(10);
3060 manifest.set_prompt(5);
3061
3062 let bytes = manifest.as_bytes();
3063 let parsed = DcpClient::parse_capability_manifest(bytes).unwrap();
3064
3065 assert!(parsed.has_tool(42));
3066 assert!(parsed.has_resource(10));
3067 assert!(parsed.has_prompt(5));
3068 }
3069
3070 #[test]
3071 fn test_capability_manifest_serialization() {
3072 let mut client = DcpClient::with_defaults();
3073 client.register_tool("tool1", "Tool 1", [0; 32], 0);
3074 client.register_resource(10);
3075 client.register_prompt(5);
3076
3077 let bytes = client.serialize_capabilities();
3078 let parsed = CapabilityManifest::from_bytes(&bytes).unwrap();
3079
3080 assert!(parsed.has_tool(1)); assert!(parsed.has_resource(10));
3082 assert!(parsed.has_prompt(5));
3083 }
3084
3085 #[test]
3086 fn test_resource_registration() {
3087 let mut client = DcpClient::with_defaults();
3088
3089 assert!(!client.has_resource(42));
3090 client.register_resource(42);
3091 assert!(client.has_resource(42));
3092 client.unregister_resource(42);
3093 assert!(!client.has_resource(42));
3094 }
3095
3096 #[test]
3097 fn test_prompt_registration() {
3098 let mut client = DcpClient::with_defaults();
3099
3100 assert!(!client.has_prompt(10));
3101 client.register_prompt(10);
3102 assert!(client.has_prompt(10));
3103 client.unregister_prompt(10);
3104 assert!(!client.has_prompt(10));
3105 }
3106
3107 #[test]
3108 fn test_extension_flags() {
3109 let mut client = DcpClient::with_defaults();
3110
3111 assert!(!client.has_extension(5));
3112 client.set_extension(5);
3113 assert!(client.has_extension(5));
3114 client.clear_extension(5);
3115 assert!(!client.has_extension(5));
3116 }
3117
3118 #[test]
3119 fn test_enforce_tool_success() {
3120 let mut client = DcpClient::with_defaults();
3121 let tool_id = client.register_tool("test", "Test", [0; 32], 0);
3122
3123 assert!(client.enforce_tool(tool_id).is_ok());
3124 }
3125
3126 #[test]
3127 fn test_enforce_tool_failure() {
3128 let client = DcpClient::with_defaults();
3129
3130 assert!(client.enforce_tool(999).is_err());
3132 }
3133
3134 #[test]
3135 fn test_enforce_resource_success() {
3136 let mut client = DcpClient::with_defaults();
3137 client.register_resource(42);
3138
3139 assert!(client.enforce_resource(42).is_ok());
3140 }
3141
3142 #[test]
3143 fn test_enforce_resource_failure() {
3144 let client = DcpClient::with_defaults();
3145
3146 assert!(client.enforce_resource(999).is_err());
3147 }
3148
3149 #[test]
3150 fn test_enforce_prompt_success() {
3151 let mut client = DcpClient::with_defaults();
3152 client.register_prompt(10);
3153
3154 assert!(client.enforce_prompt(10).is_ok());
3155 }
3156
3157 #[test]
3158 fn test_enforce_prompt_failure() {
3159 let client = DcpClient::with_defaults();
3160
3161 assert!(client.enforce_prompt(999).is_err());
3162 }
3163
3164 #[test]
3165 fn test_negotiate_capabilities() {
3166 let mut client = DcpClient::with_defaults();
3167 client.register_tool("tool1", "Tool 1", [0; 32], 0);
3168 client.register_tool("tool2", "Tool 2", [0; 32], 0);
3169 client.register_tool("tool3", "Tool 3", [0; 32], 0);
3170 client.register_resource(1);
3171 client.register_resource(2);
3172
3173 let mut server_manifest = CapabilityManifest::new(1);
3175 server_manifest.set_tool(2);
3176 server_manifest.set_tool(3);
3177 server_manifest.set_tool(4);
3178 server_manifest.set_resource(2);
3179 server_manifest.set_resource(3);
3180
3181 let negotiated = client.negotiate_capabilities(&server_manifest);
3182
3183 assert!(!negotiated.has_tool(1)); assert!(negotiated.has_tool(2)); assert!(negotiated.has_tool(3)); assert!(!negotiated.has_tool(4)); assert!(!negotiated.has_resource(1)); assert!(negotiated.has_resource(2)); assert!(!negotiated.has_resource(3)); }
3193
3194 #[test]
3195 fn test_capability_stats() {
3196 let mut client = DcpClient::with_defaults();
3197 client.register_tool("tool1", "Tool 1", [0; 32], 0);
3198 client.register_tool("tool2", "Tool 2", [0; 32], 0);
3199 client.register_resource(1);
3200 client.register_prompt(1);
3201 client.register_prompt(2);
3202 client.set_extension(0);
3203
3204 let stats = client.capability_stats();
3205
3206 assert_eq!(stats.tool_count, 2);
3207 assert_eq!(stats.resource_count, 1);
3208 assert_eq!(stats.prompt_count, 2);
3209 assert_eq!(stats.extension_count, 1);
3210 assert_eq!(stats.version, 1);
3211 }
3212
3213 #[test]
3214 fn test_capability_iterators() {
3215 let mut client = DcpClient::with_defaults();
3216 client.register_tool("tool1", "Tool 1", [0; 32], 0);
3217 client.register_tool("tool2", "Tool 2", [0; 32], 0);
3218 client.register_resource(10);
3219 client.register_resource(20);
3220 client.register_prompt(5);
3221
3222 let tool_ids: Vec<_> = client.tool_ids().collect();
3223 assert_eq!(tool_ids, vec![1, 2]);
3224
3225 let resource_ids: Vec<_> = client.resource_ids().collect();
3226 assert_eq!(resource_ids, vec![10, 20]);
3227
3228 let prompt_ids: Vec<_> = client.prompt_ids().collect();
3229 assert_eq!(prompt_ids, vec![5]);
3230 }
3231}