1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
//! Application state — a single `Clone`-able struct injected into every handler
//! via `axum::extract::State`.
//!
//! All fields are `Arc<…>` so `AppState::clone()` is cheap. Axum clones the
//! state once per request.
use std::sync::Arc;
#[cfg(feature = "telemetry")]
use cognee_observability::TelemetryGuard;
use cognee_core::PipelineRunRegistry;
use cognee_core::pipeline_run_registry::DefaultPipelineRunRegistry;
use cognee_database::{
DatabaseConnection, NoopPipelineRunRepository, PipelineRunRepository,
SeaOrmPipelineRunRepository,
};
use crate::{
auth_resolver::AuthResolver,
components::ComponentHandles,
config::{HttpServerConfig, RegistryConfig},
error::ServerError,
observability::{BufferConfig, SpanBuffer},
sync::SyncRegistry,
};
// ─── AppState ────────────────────────────────────────────────────────────────
/// Per-server dependency container shared across all handlers.
///
/// Fields that depend on closed-side features (auth chain, mailer, etc.)
/// are kept as injection seams (`Option<Arc<dyn ...>>`) so the closed
/// `cognee-http-cloud` crate can populate them via the `RouterBuilder`.
#[derive(Clone)]
pub struct AppState {
/// HTTP server config (host, port, CORS, JWT, …).
pub config: Arc<HttpServerConfig>,
/// Background pipeline-run lifecycle registry.
///
/// The inner `Arc<dyn PipelineRunRegistry>` is `Clone`-able cheaply.
pub pipelines: Arc<dyn PipelineRunRegistry>,
/// Pre-built component handles (database, storage, delete_service,
/// ontology_manager). `None` until `AppState::build` fully initialises
/// the backends — most tests leave this `None` and stub out the
/// relevant functionality directly.
pub lib: Option<Arc<ComponentHandles>>,
/// Closed-side authentication chain. `None` in pure-OSS builds; closed
/// embedders install one via `RouterBuilder::with_auth_resolver(...)`
/// or `RouterBuilder::with_extra_validator(...)`. When `None`, the
/// `AuthenticatedUser` extractor falls through to either a synthetic
/// default user (`require_authentication=false`) or a 401
/// (`require_authentication=true`).
pub auth_resolver: Option<Arc<dyn AuthResolver>>,
/// Health checker for /health endpoints. `None` falls back to a
/// synthetic `MockHealthChecker`. Embedders populate by calling
/// [`AppState::install_real_health_checker`] after wiring `lib`.
pub health: Option<Arc<dyn crate::routers::health::HealthChecker>>,
/// In-memory span buffer feeding `GET /api/v1/activity/spans`.
/// Always populated — `BufferConfig::from_env()` reads the cap. To
/// effectively disable the buffer pass `BufferConfig { max_traces: 0, .. }`.
pub spans: Arc<SpanBuffer>,
/// In-memory registry tracking one running cloud sync per user. Always
/// populated; the registry itself starts empty.
pub sync: Arc<SyncRegistry>,
/// Flush-on-drop guard for the OpenTelemetry exporter (decision 9).
/// Held only for its `Drop` side effect: the last `Arc` released calls
/// `provider.force_flush()` + `provider.shutdown()`. `None` when built
/// without explicit telemetry init (test paths, library embedders).
#[cfg(feature = "telemetry")]
pub telemetry_guard: Option<Arc<TelemetryGuard>>,
}
impl AppState {
/// Build a no-op `Arc<dyn PipelineRunRegistry>` backed by a
/// `NoopPipelineRunRepository`. Useful in tests that construct `AppState`
/// directly without a real database.
pub fn noop_pipelines() -> Arc<dyn PipelineRunRegistry> {
let repo = NoopPipelineRunRepository::arc();
let cfg = RegistryConfig::default();
DefaultPipelineRunRegistry::new(repo, cfg)
}
/// Construct an `AppState` with the given config; all optional components
/// default to `None`. Later phases call this and then set individual fields.
///
/// Builds `DefaultPipelineRunRegistry` from the config's registry knobs and
/// runs the startup orphan-reset per pipelines.md §12 — any `INITIATED` /
/// `STARTED` rows left over from a previous unclean shutdown are rewritten to
/// `ERRORED` with `reason = "server_restart_orphan"`.
pub async fn build(config: HttpServerConfig) -> Result<Self, ServerError> {
// Build an in-memory-only pipeline run repository backed by a temporary
// SQLite database. The real repository (backed by the server's own DB)
// is wired when `lib` is populated. For now we use the shared
// `cognee_database::NoopPipelineRunRepository` (gap 08-07) so the
// registry is always non-None.
let repo = NoopPipelineRunRepository::arc();
let registry_cfg = config.to_registry_config();
let pipelines: Arc<dyn PipelineRunRegistry> =
DefaultPipelineRunRegistry::new(repo, registry_cfg);
Ok(Self {
config: Arc::new(config),
pipelines,
lib: None,
auth_resolver: None,
health: None,
spans: Arc::new(SpanBuffer::new(BufferConfig::from_env())),
sync: Arc::new(SyncRegistry::new()),
#[cfg(feature = "telemetry")]
telemetry_guard: None,
})
}
/// Convenience accessor for the component handles.
///
/// Returns `None` when the server is running in test mode without backends
/// wired. Most integration tests build their own `ComponentHandles` directly.
pub fn components(&self) -> Option<&ComponentHandles> {
self.lib.as_deref()
}
/// Replace the `health` field with a `RealHealthChecker` built from the
/// currently-wired `ComponentHandles`. No-op when `lib` is `None`.
pub fn install_real_health_checker(&mut self) {
if let Some(handles) = &self.lib {
let checker = crate::health::RealHealthChecker::new(Arc::clone(handles), &self.config);
self.health = Some(Arc::new(checker));
}
}
}
// ─── Build state with a real database ─────────────────────────────────────────
impl AppState {
/// Build `AppState` with a real `DatabaseConnection` wired into the pipeline
/// registry. Used by the server startup path when backend env vars are
/// present.
///
/// Runs the orphan-reset once on startup per pipelines.md §12.
pub async fn build_with_db(
config: HttpServerConfig,
db: Arc<DatabaseConnection>,
) -> Result<Self, ServerError> {
let repo = Arc::new(SeaOrmPipelineRunRepository::new(Arc::clone(&db)))
as Arc<dyn PipelineRunRepository>;
let registry_cfg = config.to_registry_config();
// Run orphan reset on startup (best-effort — non-fatal).
let pipelines: Arc<dyn PipelineRunRegistry> =
match DefaultPipelineRunRegistry::new_with_orphan_reset(repo, registry_cfg).await {
Ok(r) => r,
Err(e) => {
tracing::warn!(
"pipeline registry startup orphan-reset failed (non-fatal): {e}"
);
// Fall back to plain new() without reset.
let repo2 = Arc::new(SeaOrmPipelineRunRepository::new(Arc::clone(&db)))
as Arc<dyn PipelineRunRepository>;
DefaultPipelineRunRegistry::new(repo2, config.to_registry_config())
}
};
Ok(Self {
config: Arc::new(config),
pipelines,
lib: None,
auth_resolver: None,
health: None,
spans: Arc::new(SpanBuffer::new(BufferConfig::from_env())),
sync: Arc::new(SyncRegistry::new()),
#[cfg(feature = "telemetry")]
telemetry_guard: None,
})
}
}