1use crate::llm::actions::{
2 client_trait::{Client, ClientActionResult},
3 protocol_trait::Protocol,
4 ActionDefinition, ParameterDefinition,
5};
6use crate::protocol::{ConnectContext, EventType};
7use crate::server::thrift::actions::{action, parameter, EXAMPLE_IDL};
8use anyhow::{bail, ensure, Result};
9use serde_json::{json, Value};
10use std::sync::LazyLock;
11
12#[derive(Default)]
13pub struct ThriftClientProtocol;
14impl ThriftClientProtocol {
15 pub fn new() -> Self {
16 Self
17 }
18}
19
20fn call() -> ActionDefinition {
21 action(
22 "thrift_call",
23 "Call a function of the IDL's service with arguments by name; Rust encodes them by their declared types and reports the result, declared exception or application error",
24 vec![
25 parameter("method", "string", "The function name, e.g. get_user", true),
26 parameter("args", "object", "Arguments by name as JSON (structs as objects, enums by label), e.g. {\"id\": 7}", false),
27 ],
28 json!({"type": "thrift_call", "method": "get_user", "args": {"id": 7}}),
29 )
30}
31fn disconnect() -> ActionDefinition {
32 action(
33 "disconnect",
34 "Close the connection",
35 vec![],
36 json!({"type": "disconnect"}),
37 )
38}
39
40pub static CONNECTED_EVENT: LazyLock<EventType> = LazyLock::new(|| {
41 EventType::new(
42 "thrift_connected",
43 "Connected; the service's functions are listed",
44 call().example.clone(),
45 )
46 .with_parameters(vec![
47 parameter("service", "string", "The service being called", true),
48 parameter(
49 "functions",
50 "array",
51 "Its functions: [{name, args, returns, oneway}]",
52 true,
53 ),
54 ])
55 .with_actions(vec![call(), disconnect()])
56});
57pub static RESULT_EVENT: LazyLock<EventType> = LazyLock::new(|| {
58 EventType::new(
59 "thrift_result",
60 "The server's answer to a call (none for oneway functions)",
61 disconnect().example.clone(),
62 )
63 .with_parameters(vec![
64 parameter("method", "string", "The function called", true),
65 parameter(
66 "result",
67 "any",
68 "The decoded return value (null for void)",
69 false,
70 ),
71 parameter(
72 "exception",
73 "object",
74 "A declared exception: {name, value}",
75 false,
76 ),
77 parameter(
78 "application_error",
79 "object",
80 "A TApplicationException: {type, message}",
81 false,
82 ),
83 ])
84 .with_actions(vec![call(), disconnect()])
85});
86
87impl Protocol for ThriftClientProtocol {
88 fn protocol_name(&self) -> &'static str {
89 "Thrift"
90 }
91 fn stack_name(&self) -> &'static str {
92 "ETH>IP>TCP>Thrift"
93 }
94 fn keywords(&self) -> Vec<&'static str> {
95 vec!["thrift", "apache thrift", "rpc client", "idl"]
96 }
97 fn description(&self) -> &'static str {
98 "Apache Thrift RPC client for a service declared in IDL, framed or unframed, binary or compact"
99 }
100 fn get_async_actions(&self, _: &crate::state::app_state::AppState) -> Vec<ActionDefinition> {
101 vec![call(), disconnect()]
102 }
103 fn get_sync_actions(&self) -> Vec<ActionDefinition> {
104 vec![]
105 }
106 fn get_event_types(&self) -> Vec<EventType> {
107 vec![CONNECTED_EVENT.clone(), RESULT_EVENT.clone()]
108 }
109 fn get_startup_parameters(&self) -> Vec<ParameterDefinition> {
110 let p = |name: &str,
111 kind: &str,
112 description: &str,
113 required: bool,
114 example: Value,
115 default: Option<Value>| ParameterDefinition {
116 name: name.into(),
117 type_hint: kind.into(),
118 description: description.into(),
119 required,
120 example,
121 default,
122 };
123 vec![
124 p(
125 "idl",
126 "string",
127 "The Thrift IDL declaring the service (no include), up to 256 KiB",
128 true,
129 json!(EXAMPLE_IDL),
130 None,
131 ),
132 p(
133 "service",
134 "string",
135 "Which service in the IDL to call; the last one when omitted",
136 false,
137 json!("Users"),
138 None,
139 ),
140 p(
141 "protocol",
142 "string",
143 "binary or compact",
144 false,
145 json!("compact"),
146 Some(json!(super::DEFAULT_PROTOCOL)),
147 ),
148 p(
149 "transport",
150 "string",
151 "framed or buffered (unframed)",
152 false,
153 json!("buffered"),
154 Some(json!(super::DEFAULT_TRANSPORT)),
155 ),
156 ]
157 }
158 fn metadata(&self) -> crate::protocol::metadata::ProtocolMetadataV2 {
159 use crate::protocol::metadata::*;
160 ProtocolMetadataV2::builder()
161 .state(DevelopmentState::Experimental)
162 .privilege_requirement(PrivilegeRequirement::None)
163 .implementation("The server's IDL parser and codecs as a client")
164 .llm_control("Which functions to call with which arguments")
165 .e2e_testing("tests/client/thrift: a thriftpy2 0.7.1 server (independent) answers calls in framed binary and buffered compact, including a declared exception and an unknown method")
166 .notes("One outstanding call at a time; no multiplexed protocol.")
167 .max_inbound_bytes(crate::server::thrift::codec::MAX_MESSAGE)
168 .build()
169 }
170 fn example_prompt(&self) -> &'static str {
171 "Call get_user(7) on the Thrift service at 127.0.0.1:9090 using this IDL"
172 }
173 fn get_startup_examples(&self) -> crate::llm::actions::StartupExamples {
174 let llm = json!({"type":"open_client","protocol":"thrift","remote_addr":"127.0.0.1:9090","instruction":"Look up user 7","startup_params":{"idl": EXAMPLE_IDL}});
175 let mut static_example = llm.clone();
176 static_example["event_handlers"] = json!([
177 {"event_pattern":"thrift_connected","handler":{"type":"static","actions":[call().example]}},
178 {"event_pattern":"thrift_result","handler":{"type":"static","actions":[{"type":"disconnect"}]}}
179 ]);
180 let mut scripted = static_example.clone();
181 scripted["event_handlers"][1]["handler"] = json!({"type":"script","language":"python","code":"import json,sys\ne=json.load(sys.stdin)['event']\nprint(json.dumps({'actions':[{'type':'disconnect'}]}))"});
182 crate::llm::actions::StartupExamples::new(llm, scripted, static_example)
183 }
184 fn group_name(&self) -> &'static str {
185 "Application"
186 }
187}
188
189impl Client for ThriftClientProtocol {
190 fn connect(
191 &self,
192 ctx: ConnectContext,
193 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<std::net::SocketAddr>> + Send>>
194 {
195 Box::pin(super::connect(ctx))
196 }
197 fn execute_action(&self, v: Value) -> Result<ClientActionResult> {
198 ensure!(
199 crate::utils::json_budget::within_budget(&v, 4 * 1024 * 1024, 200_000, 64),
200 "action exceeds the Thrift bounds"
201 );
202 match v["type"].as_str() {
203 Some("thrift_call") => {
204 ensure!(
205 v["method"]
206 .as_str()
207 .is_some_and(|m| !m.is_empty() && m.len() <= 256),
208 "method names a function"
209 );
210 if let Some(a) = v.get("args").filter(|a| !a.is_null()) {
211 ensure!(a.is_object(), "args is an object of arguments by name");
212 }
213 }
214 Some("disconnect") => return Ok(ClientActionResult::Disconnect),
215 _ => bail!("Unknown Thrift client action"),
216 }
217 Ok(ClientActionResult::Custom {
218 name: v["type"].as_str().unwrap_or_default().into(),
219 data: v,
220 })
221 }
222}