feat: support agentic ai + gaurd rails #72
Open
amitgilad3 wants to merge 4 commits into
Open
Conversation
|
Caution Review failedFailed to post review comments 📝 WalkthroughWalkthroughThis PR introduces guardrails—a rule-based safety framework for query execution—and agentic context tracking. It adds a new guardrails crate with built-in guards and chain execution, extends sessions and queries to carry agent/conversation metadata, integrates guard checks into the dispatch pipeline, and provides admin APIs and frontend UI for managing guardrails and browsing agent/conversation history. Database migrations support the new query metadata and guardrails configuration storage. ChangesGuardrails and Agentic Context System
Sequence Diagram(s)sequenceDiagram
participant Client
participant Frontend
participant Dispatch
participant GuardChain
participant Engine
participant Metrics
Client->>Frontend: Query + Agent Headers<br/>(agent-id, conversation-id, etc.)
Frontend->>Frontend: Parse headers → SessionContext<br/>(resolve agent_context)
Frontend->>Dispatch: dispatch_query(SessionContext, SQL)
Dispatch->>Dispatch: Build QueryContext<br/>(copy agent_context)
Dispatch->>GuardChain: run(GuardContext, InputLayer)<br/>(raw SQL, engine, agent info)
GuardChain->>GuardChain: Execute guards sequentially<br/>(ReadOnly, RowLimit, RequirePredicate)
alt Guard Denies
GuardChain->>Dispatch: (deny_actions, blocked=true)
Dispatch->>Metrics: Record Failed<br/>(guard_actions, was_guard_blocked)
Dispatch->>Client: Error (guard denied)
else Guards Allow/Warn
GuardChain->>Dispatch: (allow/warn_actions, blocked=false)
Dispatch->>Engine: Submit Query<br/>(with guard_actions in context)
Engine->>GuardChain: run(GuardContext, PlanLayer)
GuardChain->>Engine: (actions, blocked)
Engine->>GuardChain: run(GuardContext, OutputLayer)
GuardChain->>Engine: (actions, blocked)
Engine->>Dispatch: QueryOutcome<br/>(rows, errors)
Dispatch->>Metrics: Record QueryRecord<br/>(agent_id, conversation_id,<br/>query_intent, guard_actions,<br/>was_guard_blocked)
Dispatch->>Client: Result
end
Metrics->>Metrics: Index by agent, conversation
sequenceDiagram
participant Admin as Studio Admin
participant API as /admin/agents
participant Store as Postgres Store
participant UI as Agent Page
Admin->>API: GET /admin/agents?limit=50&offset=0
API->>Store: list_agents(50, 0)
Store->>Store: Aggregate query_count,<br/>conversation_count per agent<br/>(via query_records)
Store-->>API: Vec\<AgentSummary\>
API-->>Admin: JSON response
Admin->>UI: Display agent table<br/>(agent_id, query count,<br/>conversation count)
Admin->>UI: Click agent row
UI->>API: GET /admin/conversations?agent_id=X
API->>Store: list_conversations(agent_id=X)
Store-->>API: Vec\<ConversationSummary\>
API-->>UI: JSON response
UI->>UI: Render conversation list<br/>(with step counts, blocked flag)
Admin->>UI: Expand conversation
UI->>API: GET /admin/conversations/{id}
API->>Store: get_conversation(id)
Store-->>API: Vec\<QuerySummary\><br/>(includes guard_actions, agent context)
API-->>UI: JSON response
UI->>UI: Timeline view<br/>(guard badges, status, intent)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Poem
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
|
Contributor
There was a problem hiding this comment.
Pull request overview
Adds end-to-end “agentic” support to QueryFlux by persisting agent conversation metadata with query history and introducing a guardrails system (configurable via Studio and stored in persistence) to apply SQL safety checks before execution.
Changes:
- Persist agent context fields (agent_id, conversation_id, step_index, tool_call_id, query_intent) and guardrail audit fields (guard_actions, was_guard_blocked) to query history, plus new Admin API endpoints to browse agents and conversations.
- Introduce a new
queryflux-guardrailscrate with built-in guards and wiring into dispatch paths, along with a Studio Guardrails UI/editor. - Add documentation + sidebar updates for agent context and guardrails.
Reviewed changes
Copilot reviewed 60 out of 61 changed files in this pull request and generated 13 comments.
Show a summary per file
| File | Description |
|---|---|
| website/sidebars.ts | Adds sidebar entries for Agentic AI docs and Guardrails docs. |
| website/docs/architecture/guardrails.md | New Guardrails documentation page. |
| website/docs/agentic/agent-context.md | New Agentic context documentation page. |
| queryflux-studio/lib/script-templates.ts | Adds starter template for guard scripts. |
| queryflux-studio/lib/api.ts | Adds Studio client functions for agents, conversations, and guardrails config endpoints. |
| queryflux-studio/lib/api-types.ts | Adds types for guardrails config and agent/conversation summaries; extends query history record fields. |
| queryflux-studio/components/user-script-editor-dialog.tsx | Generalizes script editor to support script “kinds” via metadata. |
| queryflux-studio/components/guard-list.tsx | New UI component for editing an ordered list of guards (built-in/webhook/script). |
| queryflux-studio/components/guard-actions-list.tsx | New UI component for rendering guard action audit trails and summary badges. |
| queryflux-studio/components/group-form-dialog.tsx | Adds per-group guard editing (writes to guardrails config). |
| queryflux-studio/app/scripts/page.tsx | Refactors scripts page copy/metadata handling. |
| queryflux-studio/app/queries/query-table.tsx | Adds Guard column and agent pill in query history table. |
| queryflux-studio/app/queries/query-detail.tsx | Adds agent identity and guard action sections; extracts reusable detail content. |
| queryflux-studio/app/queries/page.tsx | UI polish for query history header/pagination. |
| queryflux-studio/app/guardrails/page.tsx | New Guardrails page route. |
| queryflux-studio/app/guardrails/guardrails-editor.tsx | New guardrails editor + guard script library editor. |
| queryflux-studio/app/conversations/page.tsx | New conversations listing page. |
| queryflux-studio/app/conversations/conversation-table.tsx | New sortable/searchable conversations table. |
| queryflux-studio/app/conversations/[id]/page.tsx | New conversation detail timeline page. |
| queryflux-studio/app/conversations/[id]/conversation-step-detail.tsx | Expandable per-step details for conversations. |
| queryflux-studio/app/client-shell.tsx | Adds navigation links for Agents and Guardrails. |
| queryflux-studio/app/agents/page.tsx | New agents activity page. |
| queryflux-studio/app/agents/agent-table.tsx | New sortable/searchable agents table. |
| queryflux-studio/app/agents/agent-detail.tsx | New agent detail side panel (conversations preview). |
| queryflux-studio/app/agents/[id]/page.tsx | New per-agent deep inspection page with stats. |
| queryflux-studio/app/agents/[id]/conversation-list.tsx | New expandable conversation cards/timeline list for an agent. |
| crates/queryflux/src/main.rs | Wires guardrails into LiveConfig; loads config from DB/YAML; adds reload handling. |
| crates/queryflux/Cargo.toml | Adds queryflux-guardrails dependency. |
| crates/queryflux-persistence/src/script_library.rs | Adds guard as a script kind. |
| crates/queryflux-persistence/src/query_history.rs | Extends query history DTOs; adds agent/conversation summary DTOs. |
| crates/queryflux-persistence/src/postgres/mod.rs | Adds agent/conversation queries; stores new query_record fields; persists guardrails config via ProxySettingsStore. |
| crates/queryflux-persistence/src/postgres/migrations/20260506000001_guardrails.sql | Adds guardrails table for stored guard config. |
| crates/queryflux-persistence/src/postgres/migrations/20260505000002_guard_script_kind.sql | Updates user_scripts kind constraint to include guard. |
| crates/queryflux-persistence/src/postgres/migrations/20260505000001_agent_context.sql | Adds agent context + guard audit columns and indexes to query_records. |
| crates/queryflux-persistence/src/metrics_store.rs | Adds GuardAction struct and stores guard audit fields on QueryRecord. |
| crates/queryflux-persistence/src/lib.rs | Exposes new types and extends QueryHistoryStore trait for agents/conversations. |
| crates/queryflux-persistence/src/in_memory.rs | Maps new query history fields in memory; stubs agent/conversation listing methods. |
| crates/queryflux-metrics/src/lib.rs | Re-exports GuardAction for callers. |
| crates/queryflux-guardrails/src/lib.rs | New guardrails crate entrypoint and GuardResult→GuardAction conversion. |
| crates/queryflux-guardrails/src/context.rs | New guardrails context and GuardResult definitions. |
| crates/queryflux-guardrails/src/config.rs | New guardrails config model (layered/global+groups). |
| crates/queryflux-guardrails/src/chain.rs | GuardChain execution logic and tests. |
| crates/queryflux-guardrails/src/built_in.rs | Built-in guards implementation + tests (read_only/row_limit/require_predicate). |
| crates/queryflux-guardrails/Cargo.toml | New crate manifest for guardrails. |
| crates/queryflux-frontend/src/trino_http/handlers.rs | Adds agent_context/guard fields to QueryContext/QueryOutcome for Trino polling path. |
| crates/queryflux-frontend/src/state.rs | Extends LiveConfig/QueryContext/QueryOutcome; records agent + guard fields into QueryRecord. |
| crates/queryflux-frontend/src/snowflake/tests.rs | Updates test LiveConfig initialization for new guard_chain field. |
| crates/queryflux-frontend/src/snowflake/sql_api/handlers.rs | Preserves request headers into SessionContext.extra; initializes agent_context. |
| crates/queryflux-frontend/src/snowflake/http/handlers/session.rs | Initializes agent_context in SessionContext. |
| crates/queryflux-frontend/src/snowflake/http/handlers/query.rs | Preserves request headers into SessionContext.extra; initializes agent_context. |
| crates/queryflux-frontend/src/postgres_wire/mod.rs | Passes startup params into SessionContext.extra and adds agent-context tests. |
| crates/queryflux-frontend/src/mysql_wire/mod.rs | Initializes agent_context; adds tests for SET-based agent context. |
| crates/queryflux-frontend/src/flight_sql/mod.rs | Initializes agent_context in SessionContext. |
| crates/queryflux-frontend/src/dispatch.rs | Runs guardrails before execution on async and sync dispatch paths; threads agent_context. |
| crates/queryflux-frontend/src/admin.rs | Adds agents/conversations endpoints and guardrails config endpoints. |
| crates/queryflux-frontend/Cargo.toml | Adds queryflux-guardrails dependency. |
| crates/queryflux-e2e-tests/src/harness.rs | Updates test LiveConfig initialization for new guard_chain field. |
| crates/queryflux-core/src/session.rs | Adds AgentContext + QueryIntent parsing/inference and session resolver. |
| crates/queryflux-core/src/config.rs | Adds YAML guardrails config structures to ProxyConfig. |
| CHANGELOG.md | Adds 0.1.3 release notes mentioning agentic context + guardrails. |
| Cargo.toml | Adds queryflux-guardrails workspace member and bumps version to 0.1.3. |
| Cargo.lock | Updates workspace crate versions and adds guardrails crate dependencies. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+1698
to
+1713
| let cfg = config.guardrails.as_ref()?; | ||
| let mut guards: Vec<Box<dyn Guard>> = Vec::new(); | ||
|
|
||
| for spec in &cfg.global { | ||
| match &spec.kind { | ||
| GuardKindConfig::BuiltIn => { | ||
| let name = spec.name.as_deref().unwrap_or(""); | ||
| match name { | ||
| "read_only" => guards.push(Box::new(ReadOnlyGuard)), | ||
| "row_limit" => guards.push(Box::new(RowLimitGuard { | ||
| max_rows: spec.max_rows, | ||
| })), | ||
| "require_predicate" => guards.push(Box::new(RequirePredicateGuard { | ||
| applies_to: spec.applies_to.clone().unwrap_or_default(), | ||
| })), | ||
| other => tracing::warn!(name = other, "Unknown built-in guard name; skipping"), |
Comment on lines
+1729
to
+1734
| /// Build a `GuardChain` from the flat JSON format stored by the Studio UI. | ||
| /// | ||
| /// The DB format mirrors `GuardrailsConfig` from the TypeScript API types: | ||
| /// `{ global: GuardSpecDto[], groups: Record<string, GuardSpecDto[]> }`. | ||
| /// Only the `global` array is used here; per-group dispatch is not yet implemented. | ||
| fn build_guard_chain_from_db_value(v: &serde_json::Value) -> Option<Arc<GuardChain>> { |
Comment on lines
+231
to
+253
| // Guard chain: runs after translation (SQL is final), before engine submission. | ||
| if let Some(chain) = &guard_chain { | ||
| let resolved_agent_ctx = session.resolved_agent_context(); | ||
| let guard_ctx = GuardContext { | ||
| sql: &original_sql, | ||
| translated_sql: &sql, | ||
| engine_type: &adapter.engine_type(), | ||
| cluster_group: &group, | ||
| user: session.user(), | ||
| agent_context: resolved_agent_ctx.as_ref(), | ||
| query_tags: &effective_tags, | ||
| }; | ||
| let (guard_actions, was_blocked) = chain.run(&guard_ctx, GuardLayer::Plan).await; | ||
| if was_blocked { | ||
| state.metrics.on_query_finished(&group.0, &cluster_name.0); | ||
| let _ = cluster_manager.release_cluster(&group, &cluster_name).await; | ||
| let deny_reason = guard_actions | ||
| .iter() | ||
| .find(|a| a.action == "deny") | ||
| .and_then(|a| a.reason.clone()) | ||
| .unwrap_or_else(|| "query blocked by guardrail".to_string()); | ||
| return Err(QueryFluxError::Engine(deny_reason)); | ||
| } |
| }, | ||
| query_tags: effective_tags, | ||
| query_params: vec![], | ||
| agent_context: session.agent_context.clone(), |
Comment on lines
+856
to
+857
| guard_actions: vec![], | ||
| was_guard_blocked: false, |
Comment on lines
+40
to
+46
| ```yaml | ||
| guardrails: | ||
| global: | ||
| plan: | ||
| - kind: built_in | ||
| name: read_only | ||
| ``` |
Comment on lines
+131
to
+139
| ```yaml | ||
| guardrails: | ||
| global: | ||
| plan: | ||
| - kind: | ||
| python_script: | ||
| script_id: 42 | ||
| timeout_ms: 2000 | ||
| ``` |
Comment on lines
+12
to
+20
| /// Convert a `GuardResult` into a `GuardAction` for audit recording. | ||
| pub fn result_to_action(guard_name: &str, result: &GuardResult) -> GuardAction { | ||
| match result { | ||
| GuardResult::Allow { metadata: _ } => GuardAction { | ||
| guard: guard_name.to_string(), | ||
| action: "allow".to_string(), | ||
| reason: None, | ||
| code: None, | ||
| }, |
| /// task can atomically swap the whole bundle on each reload tick. | ||
| pub struct LiveConfig { | ||
| pub router_chain: RouterChain, | ||
| /// Guard chain: global guards first, per-group guards appended at dispatch time. |
Comment on lines
+32
to
+36
| Return shapes: | ||
| {"action": "allow"} | ||
| {"action": "allow", "metadata": {"matched_rule": "...", "estimated_rows": "~50k"}} | ||
| {"action": "warn", "reason": "scans large table without partition filter"} | ||
| {"action": "deny", "reason": "missing WHERE on fct_events", "code": "MISSING_PREDICATE"} |
Contributor
Author
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
Contributor
Author
|
reference discussion - #51 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
New Features
Documentation