distri_types/
connections.rs

1use std::collections::HashMap;
2
3use chrono::{DateTime, Utc};
4use schemars::JsonSchema;
5use serde::{Deserialize, Serialize};
6use utoipa::ToSchema;
7use uuid::Uuid;
8
9use crate::McpClientTransport;
10
11#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, JsonSchema, ToSchema)]
12#[serde(rename_all = "snake_case")]
13pub enum ConnectionAuthType {
14    OAuth2,
15    Secret,
16}
17
18#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, JsonSchema, ToSchema)]
19#[serde(rename_all = "snake_case")]
20pub enum ConnectionStatus {
21    Connected,
22    Disconnected,
23    Expired,
24    NeedsSetup,
25    Partial,
26    Pending,
27    Error,
28}
29
30impl std::fmt::Display for ConnectionStatus {
31    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        match self {
33            ConnectionStatus::Connected => write!(f, "connected"),
34            ConnectionStatus::Disconnected => write!(f, "disconnected"),
35            ConnectionStatus::Expired => write!(f, "expired"),
36            ConnectionStatus::NeedsSetup => write!(f, "needs_setup"),
37            ConnectionStatus::Partial => write!(f, "partial"),
38            ConnectionStatus::Pending => write!(f, "pending"),
39            ConnectionStatus::Error => write!(f, "error"),
40        }
41    }
42}
43
44impl std::str::FromStr for ConnectionStatus {
45    type Err = anyhow::Error;
46
47    fn from_str(s: &str) -> Result<Self, Self::Err> {
48        match s {
49            "connected" => Ok(ConnectionStatus::Connected),
50            "disconnected" => Ok(ConnectionStatus::Disconnected),
51            "expired" => Ok(ConnectionStatus::Expired),
52            "needs_setup" => Ok(ConnectionStatus::NeedsSetup),
53            "partial" => Ok(ConnectionStatus::Partial),
54            "pending" => Ok(ConnectionStatus::Pending),
55            "error" => Ok(ConnectionStatus::Error),
56            _ => Err(anyhow::anyhow!("unknown connection status: {}", s)),
57        }
58    }
59}
60
61/// Unified auth scope applied to connections, bots, and tokens.
62///
63/// - `Public`: anonymous — no auth required. Valid on bots only; tokens never carry Public.
64/// - `Workspace`: a logged-in platform member (distri signup via OTP).
65/// - `User`: an external actor resolved via a connection (customer's end-user).
66#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, JsonSchema, ToSchema)]
67#[serde(rename_all = "snake_case")]
68pub enum AuthScope {
69    Public,
70    Workspace,
71    User,
72}
73
74impl std::fmt::Display for AuthScope {
75    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
76        match self {
77            Self::Public => write!(f, "public"),
78            Self::Workspace => write!(f, "workspace"),
79            Self::User => write!(f, "user"),
80        }
81    }
82}
83
84impl std::str::FromStr for AuthScope {
85    type Err = anyhow::Error;
86    fn from_str(s: &str) -> Result<Self, Self::Err> {
87        match s {
88            "public" => Ok(Self::Public),
89            "workspace" => Ok(Self::Workspace),
90            "user" => Ok(Self::User),
91            _ => Err(anyhow::anyhow!("unknown auth_scope: {}", s)),
92        }
93    }
94}
95
96/// One configurable field on a Custom-auth connection.
97#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, ToSchema, PartialEq)]
98pub struct CustomField {
99    /// Stable identifier used as the secret key suffix (e.g. `api_key`).
100    pub key: String,
101    /// Optional human-readable label for the UI.
102    #[serde(default, skip_serializing_if = "Option::is_none")]
103    pub label: Option<String>,
104    /// Mask the input in the UI and redact in logs.
105    #[serde(default)]
106    pub is_secret: bool,
107    /// Whether a value is required for the connection to be considered configured.
108    #[serde(default = "default_required")]
109    pub required: bool,
110}
111
112fn default_required() -> bool {
113    true
114}
115
116/// Semantic category for grouping Directory tiles, declared inline on
117/// each catalog entry (`"group": "communication"`). Kept as a free-form
118/// snake_case string so new buckets can be added by editing the catalog
119/// JSON alone — no Rust recompile, no enum drift. The UI owns display
120/// labels and ordering for known values and falls back to title-casing
121/// the raw slug for unknowns.
122pub type ProviderGroup = String;
123
124/// Full OAuth provider declaration carried inline on a Connection.
125///
126/// One canonical shape across three sources:
127///   * **Built-in catalog** (`additional_providers.json`) — seeded into
128///     the create form when the user picks a known provider tile.
129///   * **MCP discovery** (RFC 8414 + 9728) — seeded from
130///     `POST /v1/connections/oauth/discover` when the user pastes an
131///     MCP URL whose server publishes auth-server metadata.
132///   * **Admin-entered** — workspace admin types the values into the
133///     custom-connection form directly.
134///
135/// Once a Connection is created, the config is *frozen* on the row —
136/// catalog edits do not retroactively apply. Use
137/// `POST /v1/connections/{id}/resync-provider` to re-apply the catalog
138/// over an existing connection.
139///
140/// **Always USER-scoped.** Connections authorize an individual end-user's
141/// identity at the third party (Slack `xoxp-…`, Google as user X, …).
142/// Bot install flows use channel-setup paths (`channels.bot_token`),
143/// NOT this config.
144#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, ToSchema, PartialEq)]
145pub struct OAuthProviderConfig {
146    /// Stable slug (`slack`, `github`, `linear`, …) — drives the
147    /// catalog-resync lookup and the secret-store namespace prefix.
148    pub name: String,
149    /// Friendly label for the UI. Falls back to title-cased `name`.
150    #[serde(default, skip_serializing_if = "Option::is_none")]
151    pub display_name: Option<String>,
152    pub authorization_url: String,
153    pub token_url: String,
154    #[serde(default, skip_serializing_if = "Option::is_none")]
155    pub refresh_url: Option<String>,
156    /// RFC 7591 registration endpoint, when the auth server publishes one
157    /// (discovery sets this; catalog entries usually don't).
158    #[serde(default, skip_serializing_if = "Option::is_none")]
159    pub registration_endpoint: Option<String>,
160    #[serde(default)]
161    pub scopes_supported: Vec<String>,
162    #[serde(default)]
163    pub default_scopes: Vec<String>,
164    /// Extra auth-URL query params the provider always wants set
165    /// (Google's `access_type=offline`, Twitter PKCE flag, …).
166    #[serde(default)]
167    pub default_auth_params: HashMap<String, String>,
168    /// JSON Schema describing caller-overridable auth-URL params
169    /// (Slack's `team`, Microsoft's `tenant`, …). UI auto-renders inputs.
170    #[serde(default, skip_serializing_if = "Option::is_none")]
171    pub auth_params_schema: Option<serde_json::Value>,
172    #[serde(default)]
173    pub pkce_required: bool,
174    /// Env-var names that hold the platform's pre-registered OAuth client
175    /// (`SLACK_CLIENT_ID`/`_SECRET`). Both `None` ⇒ no platform creds,
176    /// BYOK required. UI offers BYOK alongside platform when present.
177    #[serde(default, skip_serializing_if = "Option::is_none")]
178    pub env_client_id: Option<String>,
179    #[serde(default, skip_serializing_if = "Option::is_none")]
180    pub env_client_secret: Option<String>,
181    #[serde(default, skip_serializing_if = "Option::is_none")]
182    pub icon_url: Option<String>,
183}
184
185/// Directory tile wire shape served by `GET /v1/connections/providers`.
186/// Tagged union on `kind` so the UI's MCP-vs-REST switching is
187/// exhaustive (no nullable `transport_url` to leak through type holes).
188///
189/// **Rest** — vanilla OAuth (Google OIDC, GitHub, Notion). Workflows
190/// use the stored token against the provider's REST API directly.
191///
192/// **Mcp** — OAuth + pinned `transport_url`. Picking this tile creates
193/// an MCP-kind Connection with `Connection.kind.mcp.transport.url`
194/// pre-pinned. The OAuth fields are still used for the consent flow.
195#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, ToSchema)]
196#[serde(tag = "kind", rename_all = "snake_case")]
197pub enum CatalogProvider {
198    Rest {
199        #[serde(flatten)]
200        oauth: OAuthProviderConfig,
201        #[serde(default, skip_serializing_if = "Option::is_none")]
202        group: Option<ProviderGroup>,
203    },
204    Mcp {
205        #[serde(flatten)]
206        oauth: OAuthProviderConfig,
207        /// Streamable-HTTP endpoint the connection's MCP transport will
208        /// be pinned to.
209        transport_url: String,
210        #[serde(default, skip_serializing_if = "Option::is_none")]
211        group: Option<ProviderGroup>,
212    },
213}
214
215impl CatalogProvider {
216    /// Borrow the OAuth bag shared by both variants.
217    pub fn oauth(&self) -> &OAuthProviderConfig {
218        match self {
219            Self::Rest { oauth, .. } | Self::Mcp { oauth, .. } => oauth,
220        }
221    }
222
223    /// Borrow the group label shared by both variants.
224    pub fn group(&self) -> Option<&ProviderGroup> {
225        match self {
226            Self::Rest { group, .. } | Self::Mcp { group, .. } => group.as_ref(),
227        }
228    }
229
230    /// Borrow the pinned MCP transport URL when the entry is MCP-flavored.
231    pub fn transport_url(&self) -> Option<&str> {
232        match self {
233            Self::Mcp { transport_url, .. } => Some(transport_url.as_str()),
234            _ => None,
235        }
236    }
237}
238
239impl OAuthProviderConfig {
240    /// Effective display name — friendly label or title-cased slug.
241    pub fn display(&self) -> String {
242        self.display_name.clone().unwrap_or_else(|| {
243            let mut chars = self.name.chars();
244            match chars.next() {
245                Some(c) => c.to_uppercase().collect::<String>() + chars.as_str(),
246                None => String::new(),
247            }
248        })
249    }
250
251    /// Project this provider config to the `AuthType::OAuth2 { ... }`
252    /// shape consumed by `OAuthHandler`. `scopes` is the per-connection
253    /// requested scopes (different from the catalog's `scopes_supported`).
254    pub fn to_auth_type(&self, scopes: Vec<String>) -> crate::auth::AuthType {
255        crate::auth::AuthType::OAuth2 {
256            flow_type: crate::auth::OAuth2FlowType::AuthorizationCode,
257            authorization_url: self.authorization_url.clone(),
258            token_url: self.token_url.clone(),
259            refresh_url: self.refresh_url.clone(),
260            scopes,
261            send_redirect_uri: true,
262        }
263    }
264}
265
266/// What kind of authn this Connection holds. Lives directly on the Connection
267/// row — auth is connection-shaped, not a shared entity.
268///
269/// Replaces the prior split `Credential.material` enum.
270#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, ToSchema, PartialEq)]
271#[serde(tag = "type", rename_all = "snake_case")]
272pub enum ConnectionAuth {
273    /// No auth needed. Used by open MCP servers and test connections.
274    None,
275    /// OAuth (platform-managed, BYOK, or DCR for MCP). The provider's full
276    /// declaration (auth URL, token URL, scopes, PKCE, env-var refs for
277    /// platform client creds, …) is carried inline so the connection is
278    /// self-sufficient: no runtime catalog lookup needed for OAuth flow.
279    Oauth {
280        provider: OAuthProviderConfig,
281        #[serde(default)]
282        scopes: Vec<String>,
283    },
284    /// User-supplied named fields (API keys, custom headers). Values live in
285    /// the `secrets` table under key `connection.<id>.<field_key>`.
286    Custom {
287        #[serde(default)]
288        fields: Vec<CustomField>,
289    },
290    /// Distri's own session token IS the auth.
291    DistriNative,
292}
293
294impl ConnectionAuth {
295    pub fn provider_name(&self) -> &str {
296        match self {
297            Self::None => "none",
298            Self::Oauth { provider, .. } => provider.name.as_str(),
299            Self::Custom { .. } => "custom",
300            Self::DistriNative => "distri",
301        }
302    }
303
304    /// Borrow the inline OAuth provider config when this is the `Oauth` variant.
305    pub fn oauth_config(&self) -> Option<&OAuthProviderConfig> {
306        match self {
307            Self::Oauth { provider, .. } => Some(provider),
308            _ => None,
309        }
310    }
311
312    pub fn is_oauth(&self) -> bool {
313        matches!(self, Self::Oauth { .. })
314    }
315
316    pub fn is_custom(&self) -> bool {
317        matches!(self, Self::Custom { .. })
318    }
319
320    pub fn is_distri_native(&self) -> bool {
321        matches!(self, Self::DistriNative)
322    }
323
324    pub fn custom_fields(&self) -> &[CustomField] {
325        match self {
326            Self::Custom { fields } => fields,
327            _ => &[],
328        }
329    }
330
331    pub fn custom_required_fields(&self) -> Vec<&CustomField> {
332        match self {
333            Self::Custom { fields } => fields.iter().filter(|f| f.required).collect(),
334            _ => vec![],
335        }
336    }
337}
338
339/// Typed OAuth/refresh token bundle stored in Redis under
340/// `connection:token:{connection_id}`.
341#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, ToSchema)]
342pub struct ConnectionToken {
343    pub access_token: String,
344    #[serde(default, skip_serializing_if = "Option::is_none")]
345    pub refresh_token: Option<String>,
346    #[serde(default, skip_serializing_if = "Option::is_none")]
347    pub expires_at: Option<DateTime<Utc>>,
348    #[serde(default = "default_token_type")]
349    pub token_type: String,
350    #[serde(default)]
351    pub scopes: Vec<String>,
352}
353
354fn default_token_type() -> String {
355    "Bearer".to_string()
356}
357
358impl ConnectionToken {
359    pub fn is_expired(&self) -> bool {
360        self.expires_at.map(|exp| exp < Utc::now()).unwrap_or(false)
361    }
362}
363
364/// What capability surface a connection exposes.
365///
366/// `Default` is the historical behavior: auth-only — the connection contributes
367/// credentials to the agent's env vars or to the outbound HTTP proxy.
368///
369/// `Mcp` makes the connection *also* a remote tool source. Agents reference it
370/// by `Connection.name` in `ToolsConfig.mcp[].server`; the executor builds an
371/// `McpClientPool` from every `kind = Mcp` connection in scope, with auth
372/// (`auth_type`) injected as the transport's bearer header at connect time.
373///
374/// Auth and capability are orthogonal: the same OAuth provider can back both a
375/// `Default` GitHub connection (REST proxy) and a `Mcp` GitHub connection
376/// (Streamable HTTP) — they're two rows that share `auth_type.provider`.
377#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, ToSchema, PartialEq)]
378#[serde(tag = "type", rename_all = "snake_case")]
379pub enum ConnectionKind {
380    /// REST / CLI proxy connection. Agents call the API documented in
381    /// `skill_content`; distri injects the credential's headers via
382    /// `/proxy/request` or the `inject_connection_env` tool.
383    Default {
384        /// Markdown skill describing the API surface. Required for custom
385        /// connections; optional for built-in OAuth providers that ship a
386        /// bundled template (the server fills in from
387        /// `connection_skill_templates/<provider>.md` when absent).
388        #[serde(default, skip_serializing_if = "Option::is_none")]
389        skill_content: Option<String>,
390    },
391    /// Remote MCP server. distri connects to `transport.url` and exposes
392    /// its tools to agents that reference this connection's name in
393    /// `tools.mcp[].server`.
394    Mcp {
395        #[serde(flatten)]
396        mcp: McpConnectionSpec,
397    },
398}
399
400impl Default for ConnectionKind {
401    fn default() -> Self {
402        Self::Default {
403            skill_content: None,
404        }
405    }
406}
407
408impl ConnectionKind {
409    pub fn is_mcp(&self) -> bool {
410        matches!(self, Self::Mcp { .. })
411    }
412    pub fn as_mcp(&self) -> Option<&McpConnectionSpec> {
413        match self {
414            Self::Mcp { mcp } => Some(mcp),
415            _ => None,
416        }
417    }
418    /// Skill markdown if this is a Default-kind connection that ships one.
419    /// MCP-kind connections always return `None` (their tool surface is
420    /// the MCP server's `tools/list`, not a skill doc).
421    pub fn skill_content(&self) -> Option<&str> {
422        match self {
423            Self::Default { skill_content } => skill_content.as_deref(),
424            Self::Mcp { .. } => None,
425        }
426    }
427    pub fn kind_str(&self) -> &'static str {
428        match self {
429            Self::Default { .. } => "default",
430            Self::Mcp { .. } => "mcp",
431        }
432    }
433}
434
435/// MCP-specific configuration carried on `ConnectionKind::Mcp`.
436///
437/// The transport is restricted to remote variants in the UI (Streamable HTTP /
438/// SSE) — stdio is intentionally not user-configurable because connections are
439/// meant to be portable across hosts. `extra_headers` are merged with the
440/// resolver-injected `Authorization` header at connect time.
441#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, ToSchema, PartialEq)]
442#[serde(rename_all = "snake_case")]
443pub struct McpConnectionSpec {
444    pub transport: McpClientTransport,
445    /// Optional human description shown in the UI list.
446    #[serde(default, skip_serializing_if = "Option::is_none")]
447    pub description: Option<String>,
448    /// Include/exclude glob patterns applied to discovered tool names.
449    #[serde(default, skip_serializing_if = "Option::is_none")]
450    pub tool_filter: Option<McpToolFilter>,
451    /// Whether the server is enabled for tool resolution.
452    #[serde(default = "default_true")]
453    pub enabled: bool,
454}
455
456fn default_true() -> bool {
457    true
458}
459
460#[derive(Debug, Clone, Default, Serialize, Deserialize, JsonSchema, ToSchema, PartialEq)]
461#[serde(rename_all = "snake_case")]
462pub struct McpToolFilter {
463    #[serde(default, skip_serializing_if = "Vec::is_empty")]
464    pub include: Vec<String>,
465    #[serde(default, skip_serializing_if = "Vec::is_empty")]
466    pub exclude: Vec<String>,
467}
468
469#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, ToSchema)]
470pub struct Connection {
471    pub id: Uuid,
472    pub workspace_id: Uuid,
473    pub skill_id: Uuid,
474    pub name: String,
475    pub status: ConnectionStatus,
476    pub config: serde_json::Value,
477    pub connected_by: Option<Uuid>,
478    pub created_at: DateTime<Utc>,
479    pub updated_at: DateTime<Utc>,
480    /// Who is allowed to use this connection. Workspace = platform members,
481    /// EndUser = external actors resolved via a handshake.
482    pub auth_scope: AuthScope,
483    /// The authentication material this connection carries. Auth is
484    /// connection-shaped — there is no separate Credential entity.
485    pub auth: ConnectionAuth,
486    /// What surface this connection exposes (auth-only vs MCP tool source).
487    #[serde(default)]
488    pub kind: ConnectionKind,
489    /// Platform-seeded connections (e.g. the `distri` connection) carry is_system=true
490    /// and are write-protected from user mutations.
491    #[serde(default)]
492    pub is_system: bool,
493}
494
495#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, ToSchema)]
496pub struct NewConnection {
497    pub workspace_id: Uuid,
498    pub skill_id: Uuid,
499    pub name: String,
500    pub status: ConnectionStatus,
501    pub config: serde_json::Value,
502    pub connected_by: Option<Uuid>,
503    pub auth_scope: AuthScope,
504    pub auth: ConnectionAuth,
505    #[serde(default)]
506    pub kind: ConnectionKind,
507    #[serde(default)]
508    pub is_system: bool,
509}
510
511impl Connection {
512    pub fn is_mcp(&self) -> bool {
513        self.kind.is_mcp()
514    }
515    pub fn mcp_spec(&self) -> Option<&McpConnectionSpec> {
516        self.kind.as_mcp()
517    }
518}
519
520/// Admin-authored HTTP probe attached to a *connection* — used by the
521/// "New Custom Connection" UI to confirm the configured URL + supplied
522/// auth fields actually reach the downstream service before save.
523///
524/// Stored in `connection.config['verify_request']`. Scope is intentionally
525/// per-connection: a connection is "this URL + this transport + this auth",
526/// and the probe tests that combination. Unrelated to bot gating (which is
527/// just "does this end-user hold valid auth for this connection";
528/// see `Bot.gate_connection_id`).
529#[derive(Debug, Clone, Deserialize)]
530pub struct VerifyRequest {
531    pub url: String,
532    pub method: String,
533    #[serde(default)]
534    pub headers: HashMap<String, String>,
535}
536
537impl Connection {
538    /// Returns the `verify_request` object from `config`, if present.
539    pub fn verify_request(&self) -> Option<VerifyRequest> {
540        self.config
541            .get("verify_request")
542            .and_then(|v| serde_json::from_value(v.clone()).ok())
543    }
544}
545
546/// Declarative reference to a connection that an agent definition requires.
547///
548/// Resolved at agent-run start: the orchestrator matches this against the
549/// workspace's `connections` table, fetches the secret (OAuth token / custom
550/// fields / distri-native session), and injects the result into
551/// `ExecutorContext.env_vars` + `dynamic_values.available_connections`.
552///
553/// Prefer `provider` (portable across workspaces) over `connection_id`.
554#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, ToSchema, Default)]
555pub struct ConnectionRequirement {
556    /// Match by provider name (preferred): "google", "slack", ...
557    /// Resolved against the Connection's `auth.provider` for `Oauth`,
558    /// the Connection's name for `Custom`, and `"distri"` for `DistriNative`.
559    #[serde(default, skip_serializing_if = "Option::is_none")]
560    pub provider: Option<String>,
561
562    /// Pin to a specific connection ID. Takes precedence over `provider`.
563    #[serde(default, skip_serializing_if = "Option::is_none")]
564    pub connection_id: Option<Uuid>,
565
566    /// Minimum OAuth scopes required. Resolution fails (when `required=true`)
567    /// or marks the requirement unmet (when `required=false`) if the connected
568    /// token doesn't cover all of these.
569    #[serde(default, skip_serializing_if = "Vec::is_empty")]
570    pub scopes: Vec<String>,
571
572    /// Env var name override. Default: `<PROVIDER>_TOKEN` for OAuth,
573    /// `<PROVIDER>_<FIELD_KEY>` for each Custom field.
574    #[serde(default, skip_serializing_if = "Option::is_none")]
575    pub env_var: Option<String>,
576
577    /// If true, the agent fails to start when this connection can't be resolved.
578    /// If false (default), the agent starts and the requirement is surfaced in
579    /// `{{available_providers}}` so the LLM can prompt the user to connect.
580    #[serde(default)]
581    pub required: bool,
582}
distri_types/connections.rs

distri_types/
connections.rs
