Skip to content

Commit 640ac70

Browse files
authored
feat: Add structured audit logging for MCP feature server (#6456)
* feat: Add structured audit logging for MCP feature server Signed-off-by: Siddhesh Khairnar <[email protected]> * feat: Enhance audit logging with OpenTelemetry support and JSON-RPC correlation Signed-off-by: Siddhesh Khairnar <[email protected]> * refactor: Audit MCP at protocol layer, fix empty principal, add emit lock Signed-off-by: Siddhesh Khairnar <[email protected]> * fix(mcp): Propagate request_id for audit logging Signed-off-by: Siddhesh Khairnar <[email protected]> * refactor(mcp): Update MCP server handler for audit logging Signed-off-by: Siddhesh Khairnar <[email protected]> * fix: Address review feedback for audit logging Signed-off-by: Siddhesh Khairnar <[email protected]> --------- Signed-off-by: Siddhesh Khairnar <[email protected]>
1 parent 367001b commit 640ac70

10 files changed

Lines changed: 1678 additions & 7 deletions

File tree

‎sdk/python/feast/audit/__init__.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
# Copyright 2025 The Feast Authors
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
Lines changed: 381 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,381 @@
1+
# Copyright 2025 The Feast Authors
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""
16+
Structured audit logging for the Feast feature server.
17+
18+
Emits JSONL audit events for MCP tool calls, REST requests, and
19+
authentication/authorization decisions. Sensitive payloads (tokens,
20+
entity rows, feature values) are never included.
21+
"""
22+
23+
import abc
24+
import contextvars
25+
import logging
26+
import sys
27+
import threading
28+
import uuid
29+
from datetime import datetime, timezone
30+
from typing import TYPE_CHECKING, Any, Dict, List, Optional
31+
32+
if TYPE_CHECKING:
33+
from feast.infra.feature_servers.base_config import AuditLoggingConfig
34+
35+
from pydantic import BaseModel, Field
36+
37+
logger = logging.getLogger(__name__)
38+
39+
# ContextVar used by the MCP audit wrapper to propagate a single request_id
40+
# into the internal REST call so that ``mcp.tools.call`` and
41+
# ``http.request`` events share the same identifier in SIEM.
42+
mcp_audit_request_id: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
43+
"mcp_audit_request_id", default=None
44+
)
45+
46+
47+
# ---------------------------------------------------------------------------
48+
# Audit event schema
49+
# ---------------------------------------------------------------------------
50+
51+
52+
class AuditPrincipal(BaseModel):
53+
username: str = ""
54+
roles: List[str] = Field(default_factory=list)
55+
auth_type: str = ""
56+
57+
58+
class AuditSource(BaseModel):
59+
ip: str = ""
60+
transport: str = ""
61+
62+
63+
class AuditAction(BaseModel):
64+
mcp_tool: str = ""
65+
path: str = ""
66+
67+
68+
class AuditResource(BaseModel):
69+
type: str = ""
70+
name: str = ""
71+
actions: List[str] = Field(default_factory=list)
72+
73+
74+
class AuditEvent(BaseModel):
75+
"""A single structured audit log entry.
76+
77+
Fields follow the schema proposed in feast-dev/feast#6452.
78+
No sensitive payloads (tokens, entity rows, feature values) are stored.
79+
When OpenTelemetry is active, ``trace_id`` and ``span_id`` are populated
80+
automatically for correlation with distributed traces.
81+
"""
82+
83+
event_type: str
84+
timestamp: str = ""
85+
request_id: str = ""
86+
trace_id: Optional[str] = None
87+
span_id: Optional[str] = None
88+
jsonrpc_id: Optional[str] = None
89+
principal: AuditPrincipal = Field(default_factory=AuditPrincipal)
90+
source: AuditSource = Field(default_factory=AuditSource)
91+
action: AuditAction = Field(default_factory=AuditAction)
92+
resource: AuditResource = Field(default_factory=AuditResource)
93+
outcome: str = ""
94+
duration_ms: Optional[float] = None
95+
detail: str = ""
96+
97+
def to_jsonl(self) -> str:
98+
return self.model_dump_json(exclude_none=True)
99+
100+
101+
# ---------------------------------------------------------------------------
102+
# Sink abstraction
103+
# ---------------------------------------------------------------------------
104+
105+
106+
class AuditSink(abc.ABC):
107+
"""Base class for audit event sinks."""
108+
109+
@abc.abstractmethod
110+
def emit(self, event: AuditEvent) -> None: ...
111+
112+
def close(self) -> None:
113+
pass
114+
115+
116+
class StdoutAuditSink(AuditSink):
117+
"""Write JSONL events to stdout (the default)."""
118+
119+
def emit(self, event: AuditEvent) -> None:
120+
sys.stdout.write(event.to_jsonl() + "\n")
121+
sys.stdout.flush()
122+
123+
124+
class FileAuditSink(AuditSink):
125+
"""Append JSONL events to a local file."""
126+
127+
def __init__(self, file_path: str) -> None:
128+
self._file_path = file_path
129+
self._fh = open(file_path, "a", buffering=1)
130+
131+
def emit(self, event: AuditEvent) -> None:
132+
self._fh.write(event.to_jsonl() + "\n")
133+
134+
def close(self) -> None:
135+
self._fh.close()
136+
137+
138+
class LoggerAuditSink(AuditSink):
139+
"""Emit audit events through Python's ``logging`` module at INFO level."""
140+
141+
def __init__(self, logger_name: str = "feast.audit") -> None:
142+
self._logger = logging.getLogger(logger_name)
143+
144+
def emit(self, event: AuditEvent) -> None:
145+
self._logger.info(event.to_jsonl())
146+
147+
148+
# ---------------------------------------------------------------------------
149+
# AuditLogger — the main entry point
150+
# ---------------------------------------------------------------------------
151+
152+
153+
class AuditLogger:
154+
"""Central audit logger that routes events to the configured sink.
155+
156+
Instantiate once during feature-server startup and share via
157+
``app.state.audit_logger``.
158+
"""
159+
160+
def __init__(
161+
self,
162+
sink: AuditSink,
163+
*,
164+
log_successful_reads: bool = True,
165+
) -> None:
166+
self._sink = sink
167+
self._log_successful_reads = log_successful_reads
168+
self._lock = threading.Lock()
169+
170+
# -- helpers -----------------------------------------------------------
171+
172+
@staticmethod
173+
def new_request_id() -> str:
174+
return str(uuid.uuid4())
175+
176+
@staticmethod
177+
def _utcnow_iso() -> str:
178+
return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
179+
180+
@staticmethod
181+
def _inject_otel_context(event: AuditEvent) -> None:
182+
"""Populate trace_id/span_id from the active OpenTelemetry span, if any."""
183+
if event.trace_id is not None:
184+
return
185+
try:
186+
from opentelemetry import trace as otel_trace
187+
188+
span = otel_trace.get_current_span()
189+
ctx = span.get_span_context()
190+
if ctx and ctx.is_valid:
191+
event.trace_id = format(ctx.trace_id, "032x")
192+
event.span_id = format(ctx.span_id, "016x")
193+
except ImportError:
194+
pass
195+
except Exception:
196+
pass
197+
198+
# MCP tool names and REST paths that correspond to read operations.
199+
_READ_TOOL_NAMES = frozenset(
200+
{
201+
"get_online_features",
202+
"retrieve_online_documents",
203+
"get_historical_features",
204+
}
205+
)
206+
_READ_ACTIONS = frozenset({"READ_ONLINE", "READ_OFFLINE"})
207+
208+
@classmethod
209+
def _is_read_event(cls, event: AuditEvent) -> bool:
210+
"""Return ``True`` when *event* represents a successful read.
211+
212+
Checks both ``resource.actions`` (populated by REST middleware) and
213+
``action.mcp_tool`` (populated by the MCP handler wrapper) so that
214+
suppression via ``log_successful_reads=False`` works for both layers.
215+
"""
216+
if event.event_type not in {"mcp.tools.call", "http.request"}:
217+
return False
218+
resource_actions = set(event.resource.actions)
219+
if resource_actions and resource_actions.issubset(cls._READ_ACTIONS):
220+
return True
221+
if event.action.mcp_tool in cls._READ_TOOL_NAMES:
222+
return True
223+
return False
224+
225+
# -- public API --------------------------------------------------------
226+
227+
def log(self, event: AuditEvent) -> None:
228+
if not event.timestamp:
229+
event.timestamp = self._utcnow_iso()
230+
if not event.request_id:
231+
event.request_id = self.new_request_id()
232+
233+
self._inject_otel_context(event)
234+
235+
if event.outcome == "success" and not self._log_successful_reads:
236+
if self._is_read_event(event):
237+
return
238+
239+
try:
240+
with self._lock:
241+
self._sink.emit(event)
242+
except Exception:
243+
logger.exception("Failed to emit audit event")
244+
245+
def log_mcp_call(
246+
self,
247+
*,
248+
request_id: str,
249+
tool_name: str,
250+
path: str = "",
251+
principal: Optional[AuditPrincipal] = None,
252+
source: Optional[AuditSource] = None,
253+
resource: Optional[AuditResource] = None,
254+
outcome: str = "success",
255+
duration_ms: Optional[float] = None,
256+
detail: str = "",
257+
) -> None:
258+
self.log(
259+
AuditEvent(
260+
event_type="mcp.tools.call",
261+
request_id=request_id,
262+
principal=principal or AuditPrincipal(),
263+
source=source or AuditSource(transport="mcp-http"),
264+
action=AuditAction(mcp_tool=tool_name, path=path),
265+
resource=resource or AuditResource(),
266+
outcome=outcome,
267+
duration_ms=duration_ms,
268+
detail=detail,
269+
)
270+
)
271+
272+
def log_http_request(
273+
self,
274+
*,
275+
request_id: str,
276+
method: str,
277+
path: str,
278+
principal: Optional[AuditPrincipal] = None,
279+
source: Optional[AuditSource] = None,
280+
resource: Optional[AuditResource] = None,
281+
outcome: str = "success",
282+
duration_ms: Optional[float] = None,
283+
status_code: int = 200,
284+
) -> None:
285+
self.log(
286+
AuditEvent(
287+
event_type="http.request",
288+
request_id=request_id,
289+
principal=principal or AuditPrincipal(),
290+
source=source or AuditSource(transport="http"),
291+
action=AuditAction(path=path),
292+
resource=resource or AuditResource(),
293+
outcome=outcome,
294+
duration_ms=duration_ms,
295+
detail=f"{method} {path} -> {status_code}",
296+
)
297+
)
298+
299+
def log_authn(
300+
self,
301+
*,
302+
request_id: str,
303+
outcome: str,
304+
principal: Optional[AuditPrincipal] = None,
305+
source: Optional[AuditSource] = None,
306+
detail: str = "",
307+
) -> None:
308+
event_type = "authn.success" if outcome == "success" else "authn.failure"
309+
self.log(
310+
AuditEvent(
311+
event_type=event_type,
312+
request_id=request_id,
313+
principal=principal or AuditPrincipal(),
314+
source=source or AuditSource(),
315+
outcome=outcome,
316+
detail=detail,
317+
)
318+
)
319+
320+
def log_authz(
321+
self,
322+
*,
323+
request_id: str,
324+
outcome: str,
325+
principal: Optional[AuditPrincipal] = None,
326+
resource: Optional[AuditResource] = None,
327+
detail: str = "",
328+
) -> None:
329+
self.log(
330+
AuditEvent(
331+
event_type="authz.decision",
332+
request_id=request_id,
333+
principal=principal or AuditPrincipal(),
334+
resource=resource or AuditResource(),
335+
outcome=outcome,
336+
detail=detail,
337+
)
338+
)
339+
340+
def close(self) -> None:
341+
self._sink.close()
342+
343+
344+
# ---------------------------------------------------------------------------
345+
# Factory
346+
# ---------------------------------------------------------------------------
347+
348+
_SINK_FACTORIES = {
349+
"stdout": lambda cfg: StdoutAuditSink(),
350+
"file": lambda cfg: FileAuditSink(cfg.get("file_path", "feast_audit.log")),
351+
"logger": lambda cfg: LoggerAuditSink(cfg.get("logger_name", "feast.audit")),
352+
}
353+
354+
355+
def create_audit_logger_from_config(
356+
audit_cfg: Optional["AuditLoggingConfig"],
357+
) -> Optional[AuditLogger]:
358+
"""Build an ``AuditLogger`` from an ``AuditLoggingConfig`` pydantic model.
359+
360+
Returns ``None`` when audit logging is disabled.
361+
"""
362+
if audit_cfg is None or not getattr(audit_cfg, "enabled", False):
363+
return None
364+
365+
sink_type = getattr(audit_cfg, "sink", "stdout")
366+
raw: Dict[str, Any] = {}
367+
if hasattr(audit_cfg, "file_path"):
368+
raw["file_path"] = audit_cfg.file_path
369+
if hasattr(audit_cfg, "logger_name"):
370+
raw["logger_name"] = audit_cfg.logger_name
371+
372+
factory = _SINK_FACTORIES.get(sink_type)
373+
if factory is None:
374+
logger.warning("Unknown audit sink %r, falling back to stdout", sink_type)
375+
factory = _SINK_FACTORIES["stdout"]
376+
377+
sink = factory(raw)
378+
return AuditLogger(
379+
sink,
380+
log_successful_reads=getattr(audit_cfg, "log_successful_reads", True),
381+
)

0 commit comments

Comments
 (0)