systemprompt_agent/
extension.rs

1//! `Extension` registration for the agent domain: schemas, migrations,
2//! dependencies.
3//!
4//! Copyright (c) systemprompt.io — Business Source License 1.1.
5//! See <https://systemprompt.io> for licensing details.
6
7use systemprompt_extension::prelude::*;
8
9#[derive(Debug, Clone, Copy, Default)]
10pub struct AgentExtension;
11
12impl Extension for AgentExtension {
13    fn metadata(&self) -> ExtensionMetadata {
14        ExtensionMetadata {
15            id: "agent",
16            name: "Agent",
17            version: env!("CARGO_PKG_VERSION"),
18        }
19    }
20
21    fn schemas(&self) -> Vec<SchemaDefinition> {
22        let mut schemas = conversation_schemas();
23        schemas.extend(artifact_schemas());
24        schemas.extend(context_schemas());
25        schemas.extend(task_tracking_schemas());
26        schemas.extend(service_schemas());
27        schemas
28    }
29
30    fn dependencies(&self) -> Vec<&'static str> {
31        vec!["users", "oauth", "mcp", "ai"]
32    }
33
34    fn migrations(&self) -> Vec<Migration> {
35        extension_migrations!()
36    }
37
38    fn cross_extension_tables(&self) -> Vec<&'static str> {
39        vec!["ai_requests"]
40    }
41}
42
43fn conversation_schemas() -> Vec<SchemaDefinition> {
44    vec![
45        SchemaDefinition::new("user_contexts", include_str!("../schema/user_contexts.sql"))
46            .with_required_columns(vec![
47                "context_id".into(),
48                "user_id".into(),
49                "created_at".into(),
50            ]),
51        SchemaDefinition::new("agent_tasks", include_str!("../schema/agent_tasks.sql"))
52            .with_required_columns(vec![
53                "task_id".into(),
54                "context_id".into(),
55                "status".into(),
56                "created_at".into(),
57            ]),
58        SchemaDefinition::new("task_messages", include_str!("../schema/task_messages.sql"))
59            .with_required_columns(vec![
60                "id".into(),
61                "task_id".into(),
62                "role".into(),
63                "created_at".into(),
64            ]),
65        SchemaDefinition::new("message_parts", include_str!("../schema/message_parts.sql"))
66            .with_required_columns(vec!["id".into(), "message_id".into(), "part_kind".into()]),
67    ]
68}
69
70fn context_schemas() -> Vec<SchemaDefinition> {
71    vec![
72        SchemaDefinition::new(
73            "context_agents",
74            include_str!("../schema/context_agents.sql"),
75        )
76        .with_required_columns(vec!["id".into(), "context_id".into(), "agent_name".into()]),
77        SchemaDefinition::new(
78            "context_notifications",
79            include_str!("../schema/context_notifications.sql"),
80        )
81        .with_required_columns(vec![
82            "id".into(),
83            "context_id".into(),
84            "notification_type".into(),
85        ]),
86    ]
87}
88
89fn task_tracking_schemas() -> Vec<SchemaDefinition> {
90    vec![
91        SchemaDefinition::new(
92            "task_push_notification_configs",
93            include_str!("../schema/task_push_notification_configs.sql"),
94        )
95        .with_required_columns(vec!["id".into(), "task_id".into()]),
96        SchemaDefinition::new(
97            "task_execution_steps",
98            include_str!("../schema/task_execution_steps.sql"),
99        )
100        .with_required_columns(vec![
101            "step_id".into(),
102            "task_id".into(),
103            "step_type".into(),
104        ]),
105    ]
106}
107
108fn artifact_schemas() -> Vec<SchemaDefinition> {
109    vec![
110        SchemaDefinition::new(
111            "task_artifacts",
112            include_str!("../schema/task_artifacts.sql"),
113        )
114        .with_required_columns(vec!["id".into(), "task_id".into(), "artifact_id".into()]),
115        SchemaDefinition::new(
116            "artifact_parts",
117            include_str!("../schema/artifact_parts.sql"),
118        )
119        .with_required_columns(vec!["id".into(), "artifact_id".into(), "part_kind".into()]),
120    ]
121}
122
123fn service_schemas() -> Vec<SchemaDefinition> {
124    vec![
125        SchemaDefinition::new("services", include_str!("../schema/services.sql"))
126            .with_required_columns(vec!["name".into(), "module_name".into(), "status".into()]),
127        SchemaDefinition::new(
128            "user_session_analytics",
129            include_str!("../schema/user_session_analytics.sql"),
130        ),
131    ]
132}
133
134register_extension!(AgentExtension);
systemprompt_agent/extension.rs

systemprompt_agent/
extension.rs
