netget/client/snowflake/
mod.rs

1//! Snowflake client: NetGet speaks Snowflake's connector REST protocol and the model decides
2//! what SQL to run. See `AGENTS.md` beside this file.
3//!
4//! Rust owns the HTTP and JSON envelope, the `Authorization: Snowflake Token="…"` header, the
5//! session and master tokens (renewal included), polling a query still in progress, decoding
6//! the result (JSON rowsets and Arrow streams, typed by `rowtype`) and every bound. The model
7//! sees typed rows and the server's error codes, never the envelope.
8//!
9//! Two tasks per client, both registered so `stop_client` ends them:
10//!
11//! * the **session** owns the tokens. It logs in before `connect` returns, then runs one action
12//!   at a time — the model's, or one injected from the dashboard or MCP — and turns the
13//!   response into an event.
14//! * the **turns** task asks the model about each event and hands its actions back to the
15//!   session. Each event carries the number of model turns that led to it; past
16//!   [`MAX_FOLLOWUP_DEPTH`] the model is not asked again.
17//!
18//! **The target is mandatory.** The client talks to `endpoint_url` exactly as given, or to its
19//! `remote_addr`; without one it refuses to start. It never derives
20//! `<account>.snowflakecomputing.com` from the account name, and it never follows a URL to
21//! another origin — a result's `chunks` (presigned cloud-storage URLs) are counted, not
22//! fetched, and a `getResultUrl` must be on the target.
23//!
24//! Fail closed: a refused login fails `connect` and the model is never asked anything; an LLM
25//! failure sends nothing; a response that cannot be read becomes a `snowflake_error` naming
26//! why, never an invented result.
27pub mod actions;
28pub mod arrow;
29pub mod result;
30
31use std::net::SocketAddr;
32use std::time::Duration;
33
34use anyhow::{anyhow, bail, ensure, Context, Result};
35use serde_json::{json, Value};
36use tokio::sync::mpsc;
37
38use crate::client::http_fetch::FetchClient;
39use crate::client::llm_budget::call_llm_for_client;
40use crate::llm::actions::client_trait::{Client, ClientActionResult};
41use crate::logging::emit::Log;
42use crate::protocol::{ConnectContext, Event, EventType};
43use crate::state::client_handles::{ClientCommand, ClientSendOutcome};
44use crate::state::{AccessLogOwner, ClientStatus};
45pub use actions::SnowflakeClientProtocol;
46
47/// How many model turns may follow from one another before the chain stops.
48pub const MAX_FOLLOWUP_DEPTH: u32 = 8;
49/// How long a login, or a statement with its polling, may take.
50pub const REQUEST_TIMEOUT: Duration = Duration::from_secs(120);
51const MAX_REQUEST_TIMEOUT_SECS: u64 = 3600;
52/// The largest response body read. A JSON or Arrow first result set is bounded by Snowflake at
53/// a few MiB before it moves to chunks.
54pub const MAX_RESPONSE_BYTES: usize = 16 * 1024 * 1024;
55pub const DEFAULT_MAX_ROWS: usize = 100;
56pub const MAX_ROWS_LIMIT: usize = 1000;
57/// The longest statement sent.
58pub const MAX_SQL_BYTES: usize = 1024 * 1024;
59/// How many times a query still in progress is polled.
60pub const MAX_POLLS: usize = 1000;
61const POLL_INTERVAL: Duration = Duration::from_millis(500);
62/// Events waiting for the model.
63const TURNS: usize = 64;
64/// The application id sent at login.
65pub const CLIENT_APP_ID: &str = "NetGet";
66
67const CODE_SESSION_EXPIRED: &str = "390112";
68const CODE_QUERY_IN_PROGRESS: &str = "333333";
69const CODE_QUERY_IN_PROGRESS_ASYNC: &str = "333334";
70
71fn user_agent() -> String {
72    format!("NetGet/{}", env!("CARGO_PKG_VERSION"))
73}
74
75/// The base URL this client talks to: `endpoint_url` as given, else `remote_addr` (a bare
76/// `host:port` meaning `http://host:port`). There is no third option.
77pub fn resolve_endpoint(remote_addr: &str, endpoint_url: Option<&str>) -> Result<String> {
78    let given = match endpoint_url.map(str::trim).filter(|u| !u.is_empty()) {
79        Some(url) => url.to_string(),
80        None => {
81            let raw = remote_addr.trim();
82            ensure!(
83                !raw.is_empty(),
84                "the Snowflake client needs a target: a remote_addr (host:port) or an \
85                 endpoint_url. Refusing to start without one; it never derives \
86                 <account>.snowflakecomputing.com from the account name"
87            );
88            if raw.contains("://") {
89                raw.to_string()
90            } else {
91                format!("http://{raw}")
92            }
93        }
94    };
95    let parsed = url::Url::parse(&given).with_context(|| format!("{given:?} is not a URL"))?;
96    ensure!(
97        matches!(parsed.scheme(), "http" | "https") && parsed.host_str().is_some(),
98        "{given:?} is not an http(s) URL with a host"
99    );
100    ensure!(
101        parsed.query().is_none() && parsed.fragment().is_none(),
102        "{given:?}: the endpoint is a base URL, without a query or fragment"
103    );
104    Ok(given.trim_end_matches('/').to_string())
105}
106
107#[cfg(not(target_arch = "wasm32"))]
108fn http_client(base: &str, timeout: Duration) -> Result<FetchClient> {
109    let builder = reqwest::Client::builder()
110        .timeout(timeout)
111        .redirect(reqwest::redirect::Policy::none())
112        .user_agent(user_agent());
113    let client = crate::llm::ollama_client::configured_for_endpoint(builder, base)
114        .build()
115        .context("building the HTTP client")?;
116    Ok(FetchClient::from_reqwest(client).with_max_body(MAX_RESPONSE_BYTES))
117}
118
119#[cfg(target_arch = "wasm32")]
120fn http_client(_base: &str, timeout: Duration) -> Result<FetchClient> {
121    Ok(FetchClient::transport(timeout)
122        .with_max_body(MAX_RESPONSE_BYTES)
123        .with_user_agent(&user_agent()))
124}
125
126/// What a login needs.
127#[derive(Debug, Clone, Default)]
128pub struct Login {
129    pub username: String,
130    pub password: String,
131    pub account: Option<String>,
132    pub database: Option<String>,
133    pub schema: Option<String>,
134    pub warehouse: Option<String>,
135    pub role: Option<String>,
136}
137
138/// One logged-in session.
139pub struct Session {
140    http: FetchClient,
141    pub base: String,
142    token: String,
143    master_token: Option<String>,
144    timeout: Duration,
145    sequence: u64,
146    pub bytes_out: usize,
147    /// What the login response said about the session, for the connected event.
148    pub info: Value,
149}
150
151/// What one action produced: the event to raise.
152pub type Done = (&'static EventType, Value);
153
154fn code_of(resp: &Value) -> Option<String> {
155    match resp.get("code") {
156        Some(Value::String(s)) if !s.is_empty() => Some(s.clone()),
157        Some(Value::Number(n)) => Some(n.to_string()),
158        _ => None,
159    }
160}
161
162fn text_of(v: Option<&Value>) -> Option<String> {
163    match v? {
164        Value::String(s) => Some(s.clone()),
165        Value::Null => None,
166        other => Some(other.to_string()),
167    }
168}
169
170fn request_id() -> String {
171    uuid::Uuid::new_v4().to_string()
172}
173
174impl Session {
175    async fn post(
176        &mut self,
177        path: &str,
178        body: &Value,
179        token: Option<&str>,
180    ) -> Result<(u16, Value)> {
181        let url = format!("{}{path}", self.base);
182        let payload = serde_json::to_vec(body)?;
183        self.bytes_out += payload.len();
184        let mut req = self
185            .http
186            .post(&url)
187            .header("Content-Type", "application/json")
188            .header("Accept", "application/snowflake")
189            .body(payload);
190        if let Some(token) = token {
191            req = req.header("Authorization", format!("Snowflake Token=\"{token}\""));
192        }
193        read(req.send().await.with_context(|| format!("POST {path}"))?).await
194    }
195
196    async fn get(&mut self, url: &str) -> Result<(u16, Value)> {
197        let req = self
198            .http
199            .get(url)
200            .header("Accept", "application/snowflake")
201            .header(
202                "Authorization",
203                format!("Snowflake Token=\"{}\"", self.token),
204            );
205        read(req.send().await.with_context(|| format!("GET {url}"))?).await
206    }
207
208    /// Log in. Any refusal is an error: the caller fails `connect`.
209    pub async fn login(base: String, login: &Login, timeout: Duration) -> Result<Session> {
210        ensure!(!login.username.is_empty(), "the username is empty");
211        let mut session = Session {
212            http: http_client(&base, timeout)?,
213            base,
214            token: String::new(),
215            master_token: None,
216            timeout,
217            sequence: 0,
218            bytes_out: 0,
219            info: Value::Null,
220        };
221        let query = {
222            let mut query = url::form_urlencoded::Serializer::new(String::new());
223            query.append_pair("request_id", &request_id());
224            for (key, value) in [
225                ("databaseName", &login.database),
226                ("schemaName", &login.schema),
227                ("warehouse", &login.warehouse),
228                ("roleName", &login.role),
229            ] {
230                if let Some(v) = value.as_deref().filter(|v| !v.is_empty()) {
231                    query.append_pair(key, v);
232                }
233            }
234            query.append_pair("request_guid", &request_id());
235            query.finish()
236        };
237        let mut data = json!({
238            "CLIENT_APP_ID": CLIENT_APP_ID,
239            "CLIENT_APP_VERSION": env!("CARGO_PKG_VERSION"),
240            "LOGIN_NAME": login.username,
241            "PASSWORD": login.password,
242            "CLIENT_ENVIRONMENT": {"APPLICATION": CLIENT_APP_ID, "OS": std::env::consts::OS},
243            "SESSION_PARAMETERS": {},
244        });
245        if let Some(account) = login.account.as_deref().filter(|a| !a.is_empty()) {
246            data["ACCOUNT_NAME"] = json!(account);
247        }
248        let path = format!("/session/v1/login-request?{query}");
249        let (status, resp) =
250            tokio::time::timeout(timeout, session.post(&path, &json!({ "data": data }), None))
251                .await
252                .map_err(|_| anyhow!("no login response within {}s", timeout.as_secs()))??;
253        if resp.get("success").and_then(Value::as_bool) != Some(true) {
254            bail!(
255                "login refused (HTTP {status}): code {} {} decision=refused",
256                code_of(&resp).unwrap_or_else(|| "none".into()),
257                crate::utils::sanitize::line_field(
258                    &text_of(resp.get("message")).unwrap_or_default()
259                )
260            );
261        }
262        let d = resp.get("data").context("a login success with no data")?;
263        session.token = d
264            .get("token")
265            .and_then(Value::as_str)
266            .filter(|t| !t.is_empty())
267            .context("a login success with no session token")?
268            .to_string();
269        session.master_token = d
270            .get("masterToken")
271            .and_then(Value::as_str)
272            .filter(|t| !t.is_empty())
273            .map(str::to_string);
274        let si = d.get("sessionInfo").cloned().unwrap_or(Value::Null);
275        session.info = json!({
276            "database": text_of(si.get("databaseName")),
277            "schema": text_of(si.get("schemaName")),
278            "warehouse": text_of(si.get("warehouseName")),
279            "role": text_of(si.get("roleName")),
280            "server_version": text_of(d.get("serverVersion")),
281            "session_id": text_of(d.get("sessionId")),
282        });
283        Ok(session)
284    }
285
286    /// `POST /session/token-request` with the master token, replacing the session token.
287    async fn renew(&mut self) -> Result<()> {
288        let master = self
289            .master_token
290            .clone()
291            .context("the server gave no master token to renew the session with")?;
292        let path = format!("/session/token-request?requestId={}", request_id());
293        let body = json!({"oldSessionToken": self.token, "requestType": "RENEW"});
294        let (_, resp) = self.post(&path, &body, Some(&master)).await?;
295        let d = resp.get("data").cloned().unwrap_or(Value::Null);
296        match d.get("sessionToken").and_then(Value::as_str) {
297            Some(t)
298                if resp.get("success").and_then(Value::as_bool) == Some(true) && !t.is_empty() =>
299            {
300                self.token = t.to_string();
301                if let Some(m) = d.get("masterToken").and_then(Value::as_str) {
302                    self.master_token = Some(m.to_string());
303                }
304                Ok(())
305            }
306            _ => bail!(
307                "the server refused to renew the session: code {} {}",
308                code_of(&resp).unwrap_or_else(|| "none".into()),
309                text_of(resp.get("message")).unwrap_or_default()
310            ),
311        }
312    }
313
314    /// `POST /session?delete=true`.
315    pub async fn logout(&mut self) -> Result<()> {
316        let token = self.token.clone();
317        let (status, resp) = tokio::time::timeout(
318            self.timeout,
319            self.post("/session?delete=true", &json!({}), Some(&token)),
320        )
321        .await
322        .map_err(|_| anyhow!("no logout response within {}s", self.timeout.as_secs()))??;
323        ensure!(
324            resp.get("success").and_then(Value::as_bool) == Some(true),
325            "logout refused (HTTP {status}): code {}",
326            code_of(&resp).unwrap_or_else(|| "none".into())
327        );
328        Ok(())
329    }
330
331    /// Run one statement. `Ok` carries the event to raise, a refusal included.
332    pub async fn execute(&mut self, sql: &str, max_rows: usize) -> Done {
333        let timeout = self.timeout;
334        match tokio::time::timeout(timeout, self.execute_inner(sql, max_rows)).await {
335            Ok(Ok(done)) => done,
336            Ok(Err(e)) => error_event(sql, None, format!("{e:#}"), None),
337            Err(_) => error_event(
338                sql,
339                None,
340                format!("no answer within {}s", timeout.as_secs()),
341                None,
342            ),
343        }
344    }
345
346    async fn execute_inner(&mut self, sql: &str, max_rows: usize) -> Result<Done> {
347        let mut renewed = false;
348        let (mut status, mut resp) = loop {
349            self.sequence += 1;
350            let submitted = crate::utils::clock::SystemTime::now()
351                .duration_since(crate::utils::clock::UNIX_EPOCH)
352                .map(|d| d.as_millis() as u64)
353                .unwrap_or(0);
354            let body = json!({
355                "sqlText": sql,
356                "asyncExec": false,
357                "sequenceId": self.sequence,
358                "querySubmissionTime": submitted,
359                "isInternal": false,
360                "describeOnly": false,
361            });
362            let path = format!(
363                "/queries/v1/query-request?requestId={}&request_guid={}",
364                request_id(),
365                request_id()
366            );
367            let token = self.token.clone();
368            let (status, resp) = self.post(&path, &body, Some(&token)).await?;
369            if code_of(&resp).as_deref() == Some(CODE_SESSION_EXPIRED) && !renewed {
370                renewed = true;
371                if let Err(e) = self.renew().await {
372                    return Ok(error_event(
373                        sql,
374                        Some(&resp),
375                        format!("the session expired and could not be renewed: {e:#}"),
376                        Some(status),
377                    ));
378                }
379                continue;
380            }
381            break (status, resp);
382        };
383
384        let mut polls = 0;
385        while matches!(
386            code_of(&resp).as_deref(),
387            Some(CODE_QUERY_IN_PROGRESS | CODE_QUERY_IN_PROGRESS_ASYNC)
388        ) {
389            polls += 1;
390            ensure!(
391                polls <= MAX_POLLS,
392                "the query was still in progress after {MAX_POLLS} polls"
393            );
394            let result_url = resp
395                .pointer("/data/getResultUrl")
396                .and_then(Value::as_str)
397                .context("a query in progress with no getResultUrl")?;
398            // Only ever on the target: a result URL naming another origin is refused.
399            let url = crate::client::http_fetch::resolve_same_origin(&self.base, result_url)?;
400            tokio::time::sleep(POLL_INTERVAL).await;
401            (status, resp) = self.get(&url).await?;
402        }
403
404        if resp.get("success").and_then(Value::as_bool) == Some(true) {
405            let data = resp.get("data").cloned().unwrap_or(Value::Null);
406            return Ok(match result::decode(&data, max_rows) {
407                Ok(mut out) => {
408                    out["sql"] = json!(sql);
409                    (&actions::QUERY_RESULT_EVENT, out)
410                }
411                Err(e) => error_event(
412                    sql,
413                    Some(&resp),
414                    format!("the result could not be read: {e:#}"),
415                    Some(status),
416                ),
417            });
418        }
419        let message = text_of(resp.get("message"))
420            .unwrap_or_else(|| format!("the server answered HTTP {status} without success"));
421        Ok(error_event(sql, Some(&resp), message, Some(status)))
422    }
423}
424
425fn error_event(sql: &str, resp: Option<&Value>, message: String, status: Option<u16>) -> Done {
426    let mut data = json!({"operation": "run_sql", "sql": sql, "message": message});
427    if let Some(resp) = resp {
428        if let Some(code) = code_of(resp) {
429            data["code"] = json!(code);
430        }
431        if let Some(state) = text_of(resp.pointer("/data/sqlState").or(resp.get("sqlState"))) {
432            data["sql_state"] = json!(state);
433        }
434        if let Some(id) = text_of(resp.pointer("/data/queryId")) {
435            data["query_id"] = json!(id);
436        }
437    }
438    if let Some(status) = status {
439        data["http_status"] = json!(status);
440    }
441    (&actions::ERROR_EVENT, data)
442}
443
444/// A response body as JSON, whatever the HTTP status: Snowflake reports failures in the body,
445/// and some servers send them with a 4xx or 5xx as well.
446async fn read(resp: crate::client::http_fetch::FetchResponse) -> Result<(u16, Value)> {
447    let status = resp.status().as_u16();
448    let body = resp.bytes().await.context("reading the response")?;
449    ensure!(
450        body.len() <= MAX_RESPONSE_BYTES,
451        "the response is {} bytes; at most {MAX_RESPONSE_BYTES} are read",
452        body.len()
453    );
454    let value: Value = serde_json::from_slice(&body).map_err(|e| {
455        anyhow!(
456            "the server answered HTTP {status} with a body that is not JSON ({e}): {}",
457            crate::utils::sanitize::line_field(&crate::utils::truncate::truncate_for_log(
458                &String::from_utf8_lossy(&body),
459                200
460            ))
461        )
462    })?;
463    ensure!(
464        value.is_object(),
465        "the server answered HTTP {status} with JSON that is not an object"
466    );
467    Ok((status, value))
468}
469
470pub async fn connect(ctx: ConnectContext) -> Result<SocketAddr> {
471    let params = ctx.startup_params.as_ref();
472    let string = |key: &str| -> Result<Option<String>> {
473        Ok(params
474            .map(|p| p.get_optional_string(key))
475            .transpose()?
476            .flatten())
477    };
478    let base = resolve_endpoint(&ctx.remote_addr, string("endpoint_url")?.as_deref())?;
479    crate::client::http_fetch::check_url(&base)?;
480    let login = Login {
481        username: string("username")?
482            .filter(|u| !u.is_empty())
483            .context("the Snowflake client needs the username startup parameter")?,
484        password: string("password")?
485            .context("the Snowflake client needs the password startup parameter")?,
486        account: string("account")?,
487        database: string("database")?,
488        schema: string("schema")?,
489        warehouse: string("warehouse")?,
490        role: string("role")?,
491    };
492    let timeout_secs = params
493        .map(|p| p.get_optional_u64("request_timeout_secs"))
494        .transpose()?
495        .flatten()
496        .unwrap_or(REQUEST_TIMEOUT.as_secs());
497    ensure!(
498        (1..=MAX_REQUEST_TIMEOUT_SECS).contains(&timeout_secs),
499        "request_timeout_secs is 1..={MAX_REQUEST_TIMEOUT_SECS}"
500    );
501    let log = Log::new(Some(&ctx.status_tx));
502    let mut session =
503        match Session::login(base.clone(), &login, Duration::from_secs(timeout_secs)).await {
504            Ok(s) => s,
505            Err(e) => {
506                log.warn(format!(
507                    "Snowflake client {}: login to {base} failed: {e:#} decision=fail_closed_login",
508                    ctx.client_id
509                ));
510                return Err(e);
511            }
512        };
513    ctx.state
514        .update_client_status(ctx.client_id, ClientStatus::Connected)
515        .await;
516    log.info(format!(
517        "Snowflake client {} logged in to {base} as {}",
518        ctx.client_id, login.username
519    ));
520    let _ = ctx.status_tx.send("__UPDATE_UI__".into());
521
522    // Registered before the connected turn can run, so the operator can reach the client
523    // while that turn is parked on a manual rule.
524    let external =
525        crate::client::command_support::register_command_channel(&ctx.state, ctx.client_id).await;
526    let (event_tx, event_rx) = mpsc::channel::<(Event, u32)>(TURNS);
527    let (internal_tx, internal) = mpsc::channel::<(Value, u32)>(TURNS);
528    let mut connected = json!({
529        "endpoint": base,
530        "user": login.username,
531        "account": login.account,
532    });
533    if let (Some(out), Some(info)) = (connected.as_object_mut(), session.info.as_object()) {
534        for (k, v) in info {
535            out.insert(k.clone(), v.clone());
536        }
537    }
538    let _ = event_tx.try_send((Event::new(&actions::CONNECTED_EVENT, connected), 0));
539    let dispatcher = tokio::spawn(run_turns(ctx.clone(), event_rx, internal_tx));
540    let dispatcher_abort = dispatcher.abort_handle();
541    ctx.state
542        .register_client_task(ctx.client_id, dispatcher)
543        .await;
544    let session_ctx = ctx.clone();
545    let task = tokio::spawn(async move {
546        let result = run_session(&mut session, &session_ctx, external, internal, event_tx).await;
547        dispatcher_abort.abort();
548        let status = match result {
549            Ok(()) => ClientStatus::Disconnected,
550            Err(e) => {
551                Log::new(Some(&session_ctx.status_tx)).warn(format!(
552                    "Snowflake client {} ended: {e:#} decision=connection_failed",
553                    session_ctx.client_id
554                ));
555                ClientStatus::Error(e.to_string())
556            }
557        };
558        session_ctx
559            .state
560            .update_client_status(session_ctx.client_id, status)
561            .await;
562        session_ctx
563            .state
564            .remove_client_handle(session_ctx.client_id)
565            .await;
566        let _ = session_ctx.status_tx.send("__UPDATE_UI__".into());
567    });
568    ctx.state.register_client_task(ctx.client_id, task).await;
569    // HTTP holds no one socket for the session; there is no local address to report.
570    Ok(SocketAddr::from(([0, 0, 0, 0], 0)))
571}
572
573fn reply(command: Option<ClientCommand>, outcome: ClientSendOutcome) {
574    if let Some(c) = command {
575        crate::client::command_support::reply(c, Ok(outcome));
576    }
577}
578
579/// Execute actions one at a time until the client logs out or is stopped.
580async fn run_session(
581    session: &mut Session,
582    ctx: &ConnectContext,
583    mut external: mpsc::Receiver<ClientCommand>,
584    mut internal: mpsc::Receiver<(Value, u32)>,
585    events: mpsc::Sender<(Event, u32)>,
586) -> Result<()> {
587    let protocol = SnowflakeClientProtocol;
588    loop {
589        let (action, depth, command) = tokio::select! {
590            c = external.recv() => match c {
591                Some(c) => {
592                    ctx.state.record_access_log(AccessLogOwner::Client(ctx.client_id.as_u32()), "Snowflake",
593                        None, "injected_action", c.action.clone(), vec![]).await;
594                    (c.action.clone(), 0, Some(c))
595                }
596                None => return Ok(()),
597            },
598            a = internal.recv() => match a {
599                Some((a, d)) => (a, d, None),
600                None => return Ok(()),
601            },
602        };
603        let (name, data) = match protocol.execute_action(action) {
604            Ok(ClientActionResult::Custom { name, data }) => (name, data),
605            Ok(_) => {
606                reply(
607                    command,
608                    ClientSendOutcome::Executed {
609                        detail: "nothing to send".into(),
610                    },
611                );
612                continue;
613            }
614            Err(e) => {
615                Log::new(Some(&ctx.status_tx)).warn(format!(
616                    "Snowflake client {}: action refused: {e:#}",
617                    ctx.client_id
618                ));
619                reply(
620                    command,
621                    ClientSendOutcome::Rejected {
622                        error: e.to_string(),
623                    },
624                );
625                continue;
626            }
627        };
628        if name == "logout" {
629            match session.logout().await {
630                Ok(()) => Log::new(Some(&ctx.status_tx)).info(format!(
631                    "Snowflake client {} logged out of {}",
632                    ctx.client_id, session.base
633                )),
634                Err(e) => Log::new(Some(&ctx.status_tx))
635                    .warn(format!("Snowflake client {}: logout: {e:#}", ctx.client_id)),
636            }
637            reply(command, ClientSendOutcome::Disconnected);
638            return Ok(());
639        }
640        let sql = data.get("sql").and_then(Value::as_str).unwrap_or_default();
641        let max_rows = data
642            .get("max_rows")
643            .and_then(Value::as_u64)
644            .unwrap_or(DEFAULT_MAX_ROWS as u64)
645            .min(MAX_ROWS_LIMIT as u64) as usize;
646        let before = session.bytes_out;
647        let (event_type, payload) = session.execute(sql, max_rows).await;
648        reply(
649            command,
650            ClientSendOutcome::Sent {
651                bytes_sent: session.bytes_out - before,
652            },
653        );
654        tracing::info!(
655            "Snowflake client {} {name} -> {} (depth {depth})",
656            ctx.client_id,
657            event_type.id
658        );
659        if events
660            .try_send((Event::new(event_type, payload), depth))
661            .is_err()
662        {
663            Log::new(Some(&ctx.status_tx)).warn(format!(
664                "Snowflake client {}: {} events are waiting for the model; this one is dropped \
665                 decision=turn_queue_full",
666                ctx.client_id, TURNS
667            ));
668        }
669        let _ = ctx.status_tx.send("__UPDATE_UI__".into());
670    }
671}
672
673/// Ask the model about each event, in order, and hand its actions to the session.
674async fn run_turns(
675    ctx: ConnectContext,
676    mut events: mpsc::Receiver<(Event, u32)>,
677    internal: mpsc::Sender<(Value, u32)>,
678) {
679    let protocol = SnowflakeClientProtocol;
680    while let Some((event, depth)) = events.recv().await {
681        if depth >= MAX_FOLLOWUP_DEPTH {
682            Log::new(Some(&ctx.status_tx)).warn(format!(
683                "Snowflake client {} not asking the model about {}: {depth} turns deep \
684                 decision=followup_depth",
685                ctx.client_id,
686                event.id()
687            ));
688            continue;
689        }
690        let instruction = ctx
691            .state
692            .get_instruction_for_client(ctx.client_id)
693            .await
694            .unwrap_or_default();
695        let memory = ctx
696            .state
697            .get_memory_for_client(ctx.client_id)
698            .await
699            .unwrap_or_default();
700        match call_llm_for_client(
701            &ctx.llm_client,
702            &ctx.state,
703            ctx.client_id.to_string(),
704            &instruction,
705            &memory,
706            Some(&event),
707            &protocol,
708            &ctx.status_tx,
709        )
710        .await
711        {
712            Ok(result) => {
713                if let Some(memory) = result.memory_updates {
714                    ctx.state.set_memory_for_client(ctx.client_id, memory).await;
715                }
716                tracing::info!(
717                    "Snowflake client {} {} decision={} actions={}",
718                    ctx.client_id,
719                    event.id(),
720                    if result.actions.is_empty() {
721                        "model_silent"
722                    } else {
723                        "model_actions"
724                    },
725                    result.actions.len()
726                );
727                for action in result.actions {
728                    if internal.send((action, depth + 1)).await.is_err() {
729                        return;
730                    }
731                }
732            }
733            // Fail closed: nothing is sent on the model's behalf when it cannot be asked.
734            Err(e) => Log::new(Some(&ctx.status_tx)).warn(format!(
735                "Snowflake client {} {}: {e:#} decision=fail_closed_llm_error",
736                ctx.client_id,
737                event.id()
738            )),
739        }
740    }
741}
netget/client/snowflake/mod.rs

netget/client/snowflake/
mod.rs
