Coverage for src/lilbee/mcp_server.py: 100%
705 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""MCP server exposing lilbee as tools for AI agents."""
3from __future__ import annotations
5import asyncio
6import concurrent.futures
7import functools
8import inspect
9import json
10import logging
11import os
12import re
13import threading
14import uuid
15from collections.abc import Callable, Iterator
16from contextlib import contextmanager
17from contextvars import ContextVar
18from copy import deepcopy
19from dataclasses import asdict
20from pathlib import Path
21from typing import TYPE_CHECKING, Any, TypeVar, cast
22from weakref import WeakKeyDictionary
24import anyio
25from mcp.server.mcpserver import Context, MCPServer
26from mcp.types import Tool as MCPTool
28from lilbee.app.memory import (
29 MEMORY_DISABLED_HINT,
30 forget,
31 list_memories,
32 memory_enabled,
33 recall,
34 remember,
35)
36from lilbee.app.placement import (
37 PlacementView,
38 get_placement,
39 placement_refused_message,
40 preview_placement,
41 set_placement,
42)
43from lilbee.app.search import clean_result
44from lilbee.app.services import get_services, reset_services, reset_store
45from lilbee.app.settings import (
46 SettingInfo,
47 apply_settings_update,
48 get_setting,
49 list_settings,
50 provider_reset_refused_message,
51 requires_services_reset,
52 reset_settings,
53)
54from lilbee.catalog.types import ModelSource
55from lilbee.core.config import cfg
56from lilbee.core.config.enums import CrawlRenderMode
57from lilbee.core.settings import overlay_persisted_settings
58from lilbee.core.system import LOCAL_ROOT_DIRNAME, canonical_data_root
59from lilbee.crawler import crawler_available, is_url, require_valid_crawl_url
60from lilbee.crawler.task import get_task, start_crawl
61from lilbee.data.store import (
62 EmbeddingModelMismatchError,
63 MemoryKind,
64 MemorySource,
65 SearchScope,
66 agent_owner,
67 scope_to_chunk_type,
68)
69from lilbee.runtime.cancellation import TaskCancelledError
70from lilbee.sessions import (
71 AGENT_SESSIONS_DISABLED_HINT,
72 MessageRole,
73 Session,
74 SessionMessage,
75 SessionNotFoundError,
76 SessionOrigin,
77 SessionOwnershipError,
78 TitleSource,
79 agent_sessions_enabled,
80)
81from lilbee.wiki.shared import (
82 INVALID_DRAFT_SLUG_ERROR,
83 WIKI_DISABLED_ERROR,
84 WikiSubdir,
85 total_wiki_pages,
86)
88if TYPE_CHECKING:
89 from lilbee.providers.fleet.placement_spec import PlacementSpec
91log = logging.getLogger(__name__)
93_INSTRUCTIONS = (
94 "Local search engine over the user's files, code, and crawled pages. "
95 "For any question about the user's own documents or codebase -- a lookup, a "
96 "find-in-docs, 'where is X', 'how does Y work here' -- call lilbee_search first "
97 "and answer from its cited chunks. Prefer it over web-fetch or file-read tools: "
98 "those cannot see the indexed corpus."
99)
102class _TransportState:
103 """Process-level MCP transport facts.
105 ``http_mounted`` is True only when MCP is exposed over the shared
106 streamable-http daemon (set by build_mcp_mount). On that transport multiple
107 agents share one process and one global cfg/Services singleton, so
108 vault-switching (init) and factory reset are refused: switching or tearing
109 down the store under concurrent in-flight handlers is a use-after-close /
110 identity race. stdio (one agent per process) keeps both.
111 """
113 http_mounted: bool = False
116_transport = _TransportState()
119def set_http_mounted(value: bool) -> None:
120 """Mark whether this process serves MCP over the shared HTTP daemon."""
121 _transport.http_mounted = value
124_F = TypeVar("_F", bound=Callable[..., Any])
126# Set by the offload for the duration of one sync handler and read by the long
127# ones. A parameter would reach the generated tool schema and invite an agent to
128# pass it; the handlers keep the plain signature that in-process callers use.
129_CANCEL: ContextVar[threading.Event | None] = ContextVar("lilbee_mcp_cancel", default=None)
132def _caller_cancelled() -> threading.Event | None:
133 """The stop token for the handler running on this thread, if the offload set one."""
134 return _CANCEL.get()
137def _offload_sync(fn: _F) -> _F:
138 """Run a sync tool handler off the event loop; async handlers pass through.
140 The bundled mcp SDK calls sync tool handlers directly on the event loop, so
141 under the shared streamable-http daemon one slow handler would stall every
142 connected agent. ``functools.wraps`` preserves the wrapped signature so the
143 generated tool schema is unchanged.
145 ``abandon_on_cancel`` releases the caller the moment it cancels. A worker
146 thread cannot be interrupted, so the handler still runs to completion; the
147 default would additionally hold the cancelling task until it did, which on
148 the shared mount keeps a connection open for a request nobody is waiting
149 for. Work that must actually stop takes a cancel token instead, as the
150 ingest tools do.
151 """
152 if inspect.iscoroutinefunction(fn):
153 return fn
155 @functools.wraps(fn)
156 async def _runner(*args: Any, **kwargs: Any) -> Any:
157 with _cancel_token() as token:
158 _CANCEL.set(token)
159 return await anyio.to_thread.run_sync(
160 functools.partial(fn, *args, **kwargs), abandon_on_cancel=True
161 )
163 return cast("_F", _runner)
166# (handler, wire name override, gate) applied by build_mcp_server per instance.
167_REGISTRATIONS: list[tuple[Callable[..., Any], str | None, Callable[[], bool] | None]] = []
170def _tool(fn: _F) -> _F:
171 """Register *fn* as an MCP tool with sync handlers offloaded off the loop.
173 Returns the original callable so in-process callers (tests, the stdio
174 fallback) keep the synchronous API while the schema sees the offloaded form.
175 """
176 _REGISTRATIONS.append((_offload_sync(fn), None, None))
177 return fn
180def _tool_named(name: str) -> Callable[[_F], _F]:
181 """Register an MCP tool under an explicit wire *name* (sync handlers offloaded)."""
183 def deco(fn: _F) -> _F:
184 _REGISTRATIONS.append((_offload_sync(fn), name, None))
185 return fn
187 return deco
190def _tool_if(when: Callable[[], bool]) -> Callable[[_F], _F]:
191 """Register an MCP tool gated on *when*, evaluated at server-build time.
193 The function stays importable so direct callers (tests, in-process
194 fallback) can still reach it. A server built after a config change
195 carries the current tool surface; live servers keep theirs.
196 """
197 if not callable(when):
198 raise TypeError("_tool_if takes a zero-arg callable, evaluated per build")
200 def deco(fn: _F) -> _F:
201 _REGISTRATIONS.append((_offload_sync(fn), None, when))
202 return fn
204 return deco
207def _wiki_enabled() -> bool:
208 return cfg.wiki
211def _error(msg: str) -> dict[str, Any]:
212 """Uniform error envelope MCP tool handlers return on a failure path.
214 Typed as ``dict[str, Any]`` rather than a TypedDict so it composes
215 with the success-side returns under the existing handler signatures
216 without forcing every caller to widen its return type.
217 """
218 return {"error": msg}
221def _wiki_tool(fn: _F) -> _F:
222 """Register a wiki MCP tool: gated at build time, re-checked per call.
224 A live server keeps the tool surface it was built with, so every handler
225 re-reads ``cfg.wiki`` and refuses once the setting is turned off at runtime.
226 """
228 @functools.wraps(fn)
229 def guarded(*args: Any, **kwargs: Any) -> Any:
230 if not cfg.wiki:
231 return _error(WIKI_DISABLED_ERROR)
232 return fn(*args, **kwargs)
234 checked = cast("_F", guarded)
235 _tool_if(_wiki_enabled)(checked)
236 return checked
239@_tool
240def search(
241 query: str, top_k: int | None = None, scope: str = SearchScope.BOTH.value
242) -> list[dict[str, Any]] | dict[str, Any]:
243 """Search the user's indexed documents, code, and crawled pages; prefer it over web-fetch or
244 file-read tools. Returns chunks with citations. ``scope``: "both" (default) / "raw" / "wiki"."""
245 if not query or not query.strip():
246 return _error("query must not be empty")
247 try:
248 chunk_type = scope_to_chunk_type(scope)
249 except ValueError:
250 # Smaller models routinely echo prose like "indexed docs" or "all"
251 # back as the scope value. Treat unrecognised scopes as the default
252 # "both" rather than a hard failure so the request still does work.
253 log.warning("lilbee_search: unknown scope %r, falling back to %r", scope, SearchScope.BOTH)
254 chunk_type = scope_to_chunk_type(SearchScope.BOTH.value)
255 if top_k is not None and top_k < 1:
256 # Same lenient stance as the scope fallback above: do the search with
257 # the configured default rather than hard-failing the agent's call.
258 log.warning("lilbee_search: top_k %d is not positive, using the default", top_k)
259 top_k = None
260 effective_top_k = top_k if top_k is not None else cfg.top_k
261 try:
262 results = get_services().searcher.search(
263 query, top_k=effective_top_k, chunk_type=chunk_type
264 )
265 return [clean_result(r) for r in results]
266 except EmbeddingModelMismatchError as exc:
267 # Structured so an agent can offer to adopt the index's embedder
268 # rather than parse prose out of a generic error. Names the embedder,
269 # which the HTTP search route keeps out of its generic 503.
270 return {
271 "error": str(exc),
272 "code": "INDEX_EMBEDDER_MISMATCH",
273 "persisted_model": exc.persisted_model,
274 "persisted_dim": exc.persisted_dim,
275 "adoptable": exc.dims_match,
276 }
277 except Exception as exc:
278 return _error(str(exc))
281@_tool
282def status() -> dict[str, Any]:
283 """Show indexed documents, configuration, and chunk counts."""
284 sources = get_services().store.get_sources()
285 return {
286 "config": {
287 "documents_dir": str(cfg.documents_dir),
288 "data_dir": str(cfg.data_dir),
289 "chat_model": cfg.chat_model,
290 "embedding_model": cfg.embedding_model,
291 "vision_model": cfg.vision_model,
292 "reranker_model": cfg.reranker_model,
293 "enable_ocr": cfg.enable_ocr,
294 "num_ctx": cfg.num_ctx,
295 "num_ctx_max": cfg.num_ctx_max,
296 "chat_n_ctx_target": cfg.chat_n_ctx_target,
297 "flash_attention": cfg.flash_attention,
298 "kv_cache_type": cfg.kv_cache_type.value,
299 "n_gpu_layers": cfg.n_gpu_layers,
300 "cpu_moe": cfg.cpu_moe,
301 "n_cpu_moe": cfg.n_cpu_moe,
302 "main_gpu": cfg.main_gpu,
303 "gpu_devices": cfg.gpu_devices,
304 },
305 "sources": [
306 {"filename": s["filename"], "chunk_count": s["chunk_count"]}
307 for s in sorted(sources, key=lambda x: x["filename"])
308 ],
309 "total_chunks": sum(s["chunk_count"] for s in sources),
310 "entities": _entity_status_dict(),
311 }
314def _entity_status_dict() -> dict[str, Any] | None:
315 """Entity types + extracted rows, mirroring the HTTP status section."""
316 from lilbee.app.status import entity_status
318 section = entity_status()
319 return section.model_dump() if section is not None else None
322@contextmanager
323def _cancel_token() -> Iterator[threading.Event]:
324 """A stop token for a long operation, set on any abnormal exit.
326 An agent cancelling a tool call gets what the HTTP surface gets on a client
327 disconnect: the work stops at its next boundary and keeps what it finished,
328 rather than running the rest of a corpus, download or wiki build for a
329 request that has gone away. The work runs on a thread asyncio cannot
330 interrupt, which is why cancellation alone does not reach it and the token
331 has to be polled.
333 Set on every exception, not only cancellation: a pass that raised is over
334 either way, and a caller vanishing does not always surface as a cancel.
335 """
336 token = threading.Event()
337 try:
338 yield token
339 except BaseException:
340 token.set()
341 raise
344@_tool
345async def sync(force_rebuild: bool = False, retry_skipped: bool = False) -> dict[str, Any]:
346 """Sync the documents directory into the vector store.
348 ``force_rebuild`` drops every table and re-ingests. ``retry_skipped``
349 clears failed-file skip markers without dropping the store.
350 """
351 from lilbee.data.ingest import sync as run_sync
353 with _cancel_token() as cancel:
354 return (
355 await run_sync(
356 quiet=True,
357 force_rebuild=force_rebuild,
358 retry_skipped=retry_skipped,
359 cancel=cancel,
360 )
361 ).model_dump()
364@_tool
365async def add(
366 paths: list[str],
367 force: bool = False,
368 enable_ocr: bool | None = None,
369 ocr_timeout: float | None = None,
370 render_mode: CrawlRenderMode | None = None,
371) -> dict[str, Any]:
372 """Add files, directories, or URLs to the knowledge base, then sync.
373 Paths must be absolute; URLs are crawled as markdown."""
374 from lilbee.app.ingest import register_sources
375 from lilbee.data.ingest import sync as run_sync
377 errors: list[str] = []
378 valid: list[Path] = []
379 urls: list[str] = []
380 for p_str in paths:
381 if is_url(p_str):
382 urls.append(p_str)
383 else:
384 p = Path(p_str)
385 if not p.exists():
386 errors.append(p_str)
387 else:
388 valid.append(p)
390 # Crawl URLs
391 crawled_count = 0
392 if urls:
393 from lilbee.crawler import crawler_available
395 if not crawler_available():
396 return _error("Web crawling requires: pip install 'lilbee[crawler]'")
397 from lilbee.crawler import crawl_and_save
399 for url in urls:
400 try:
401 # URL validation resolves the host (blocking DNS); run it off the
402 # event loop like the sibling crawl tool does.
403 await anyio.to_thread.run_sync(require_valid_crawl_url, url)
404 except ValueError as exc:
405 errors.append(f"{url}: {exc}")
406 continue
407 crawled_paths = await crawl_and_save(url, render_mode=render_mode)
408 crawled_count += len(crawled_paths)
410 # Registration touches config.toml (a locked read-modify-write); keep the
411 # blocking disk I/O off the event loop.
412 reg_result = await anyio.to_thread.run_sync(
413 functools.partial(register_sources, valid, force=force)
414 )
416 from lilbee.app.ingest import temporary_ocr_config
418 with temporary_ocr_config(enable_ocr, ocr_timeout), _cancel_token() as cancel:
419 sync_result = (await run_sync(quiet=True, cancel=cancel)).model_dump()
421 result: dict[str, Any] = {
422 "command": "add",
423 "copied": reg_result.registered,
424 "skipped": reg_result.skipped,
425 "tracked": reg_result.tracked,
426 "crawled": crawled_count,
427 "errors": errors,
428 "sync": sync_result,
429 }
430 if errors or sync_result.get("failed"):
431 result["warning"] = "some files could not be processed"
432 return result
435@_tool_if(crawler_available)
436async def crawl(
437 url: str,
438 depth: int | None = None,
439 max_pages: int | None = None,
440 render_mode: CrawlRenderMode | None = None,
441 include_subdomains: bool = False,
442) -> dict[str, Any]:
443 """Start a non-blocking crawl; poll via ``crawl_status(task_id)``.
444 ``depth=None`` = whole site, ``0`` = single URL. ``render_mode``: "http"/"browser"."""
445 from lilbee.crawler import crawler_available
447 if not crawler_available():
448 return _error("Web crawling requires: pip install 'lilbee[crawler]'")
449 # Mirror the REST CrawlRequest bounds so a negative value is a clean error,
450 # not an unbounded crawl.
451 if depth is not None and depth < 0:
452 return _error("depth must be 0 or greater (omit it to crawl the whole site)")
453 if max_pages is not None and max_pages < 0:
454 return _error("max_pages must be 0 or greater (0 = unlimited, omit for the safety cap)")
455 try:
456 # URL validation resolves the host (blocking DNS), so it runs off the loop.
457 # The crawl itself must be scheduled ON the loop: start_crawl uses
458 # asyncio.create_task, which requires a running event loop.
459 await anyio.to_thread.run_sync(require_valid_crawl_url, url)
460 except ValueError as exc:
461 return _error(str(exc))
463 task_id = start_crawl(
464 url,
465 depth=depth,
466 max_pages=max_pages,
467 render_mode=render_mode,
468 include_subdomains=include_subdomains,
469 )
470 return {"status": "started", "task_id": task_id, "url": url}
473@_tool_if(crawler_available)
474def crawl_status(task_id: str) -> dict[str, Any]:
475 """Poll a crawl task by id; returns ``{status, pages, error}``."""
476 task = get_task(task_id)
477 if task is None:
478 return _error(f"No task found with id: {task_id}")
479 return {
480 "task_id": task.task_id,
481 "url": task.url,
482 "status": task.status.value,
483 "pages_crawled": task.pages_crawled,
484 "pages_total": task.pages_total,
485 "error": task.error,
486 "started_at": task.started_at,
487 "finished_at": task.finished_at,
488 }
491@_tool
492def crawl_cancel(task_id: str) -> dict[str, Any]:
493 """Stop a running crawl started by ``crawl``. Pages already saved are kept."""
494 from lilbee.crawler.task import cancel_crawl
496 if get_task(task_id) is None:
497 return _error(f"No task found with id: {task_id}")
498 return {"command": "crawl_cancel", "task_id": task_id, "cancelling": cancel_crawl(task_id)}
501@_tool
502def init(path: str = "") -> dict[str, Any]:
503 """Initialize a local ``.lilbee/`` knowledge base; empty path = cwd.
505 Switches the MCP session to use it for subsequent calls.
506 """
507 if _transport.http_mounted:
508 return _error(
509 "init is unavailable on the HTTP server: it is bound to one vault and "
510 "shared by every connected client. Start a separate server for another vault."
511 )
512 # Canonical so this vault keys the same lock paths a CLI or server process
513 # would derive for the same directory.
514 base = canonical_data_root(path) if path else canonical_data_root(Path.cwd())
515 root = base / LOCAL_ROOT_DIRNAME
517 created = False
518 if not root.is_dir():
519 (root / "documents").mkdir(parents=True)
520 (root / "data").mkdir(parents=True)
521 (root / ".gitignore").write_text("data/\n", encoding="utf-8")
522 created = True
524 # Switch MCP session to this project's KB. Overlay any persisted
525 # config.toml so per-vault model / generation settings take effect,
526 # matching the CLI's --data-dir behaviour. Env export mirrors
527 # cli/app.py::_apply_data_root for worker-log parity.
528 cfg.data_root = base
529 cfg.documents_dir = root / "documents"
530 cfg.data_dir = root / "data"
531 cfg.lancedb_dir = root / "data" / "lancedb"
532 os.environ["LILBEE_DATA"] = str(base)
533 overlay_persisted_settings(base)
534 reset_services()
536 return {"command": "init", "path": str(root), "created": created}
539@_tool
540def remove(names: list[str]) -> dict[str, Any]:
541 """Remove documents from the index by source name, folder, or glob pattern.
543 Source files are never deleted. A folder name removes every document indexed
544 beneath it; a glob (``*``/``?``/``[]``) removes every matching source."""
545 from lilbee.app.ingest import remove_documents_durably
547 result = remove_documents_durably(names)
548 return {"command": "remove", "removed": result.removed, "not_found": result.not_found}
551@_tool
552def list_documents() -> dict[str, Any]:
553 """List all indexed documents with their chunk counts."""
554 sources = get_services().store.get_sources()
555 return {
556 "documents": [
557 {"filename": s["filename"], "chunk_count": s.get("chunk_count", 0)} for s in sources
558 ],
559 "total": len(sources),
560 }
563_AGENT_ORIGINS = frozenset({SessionOrigin.MCP})
566def _require_agent_session(session_id: str) -> Session:
567 """Human conversations are private: foreign ids answer not-found, so an
568 agent cannot even probe which ones exist."""
569 session = get_services().session_store.get(session_id)
570 if session.meta.origin is not SessionOrigin.MCP:
571 raise SessionNotFoundError(session_id)
572 return session
575@_tool_if(agent_sessions_enabled)
576def sessions_list() -> dict[str, Any]:
577 """List the agent's sessions, newest first."""
578 if not agent_sessions_enabled():
579 return _error(AGENT_SESSIONS_DISABLED_HINT)
580 metas = get_services().session_store.list(origins=_AGENT_ORIGINS)
581 return {"sessions": [asdict(meta) for meta in metas], "total": len(metas)}
584@_tool_if(agent_sessions_enabled)
585def session_get(session_id: str) -> dict[str, Any]:
586 """Return one agent session: metadata, transcript, summary."""
587 if not agent_sessions_enabled():
588 return _error(AGENT_SESSIONS_DISABLED_HINT)
589 try:
590 session = _require_agent_session(session_id)
591 except SessionNotFoundError as exc:
592 return _error(str(exc))
593 return {
594 "meta": asdict(session.meta),
595 "messages": [
596 {
597 "role": message.role.value,
598 "content": message.content,
599 "sources": list(message.sources),
600 "ts": message.ts,
601 }
602 for message in session.messages
603 ],
604 # What compaction folded the older turns into (empty if never compacted).
605 # An agent that resumes and continues the conversation needs it, or it
606 # rebuilds history without what was already condensed.
607 "summary": session.summary,
608 }
611@_tool_if(agent_sessions_enabled)
612def session_create(model_ref: str, scope: str = "both") -> dict[str, Any]:
613 """Start a saved chat session; returns its id."""
614 if not agent_sessions_enabled():
615 return _error(AGENT_SESSIONS_DISABLED_HINT)
616 session_id = get_services().session_store.create(
617 model_ref=model_ref, scope=scope, origin=SessionOrigin.MCP
618 )
619 return {"id": session_id, "model_ref": model_ref, "scope": scope}
622@_tool_if(agent_sessions_enabled)
623def session_add_message(
624 session_id: str,
625 role: MessageRole,
626 content: str,
627 sources: list[str] | None = None,
628 claim: bool = False,
629) -> dict[str, Any]:
630 """Append one turn; a foreign session errors unless claim=True (ask first)."""
631 if not agent_sessions_enabled():
632 return _error(AGENT_SESSIONS_DISABLED_HINT)
633 store = get_services().session_store
634 try:
635 # Re-coerce: the MCP layer passes the enum, but a raw string still
636 # arrives via direct library calls, and a bad one must error cleanly.
637 message = SessionMessage(
638 role=MessageRole(role), content=content, sources=tuple(sources or ())
639 )
640 if claim:
641 store.transfer(session_id, SessionOrigin.MCP)
642 store.add_message(session_id, message, surface=SessionOrigin.MCP)
643 except (SessionNotFoundError, SessionOwnershipError) as exc:
644 return _error(str(exc))
645 except ValueError as exc:
646 return _error(f"invalid role {role!r}: {exc}")
647 return {"id": session_id, "added": True}
650@_tool_if(agent_sessions_enabled)
651def session_set_summary(session_id: str, summary: str) -> dict[str, Any]:
652 """Replace an agent session's compaction summary."""
653 if not agent_sessions_enabled():
654 return _error(AGENT_SESSIONS_DISABLED_HINT)
655 try:
656 _require_agent_session(session_id)
657 get_services().session_store.set_summary(session_id, summary)
658 except SessionNotFoundError as exc:
659 return _error(str(exc))
660 return {"id": session_id, "summary": summary}
663@_tool_if(agent_sessions_enabled)
664def session_rename(session_id: str, title: str) -> dict[str, Any]:
665 """Rename an agent session."""
666 if not agent_sessions_enabled():
667 return _error(AGENT_SESSIONS_DISABLED_HINT)
668 try:
669 _require_agent_session(session_id)
670 get_services().session_store.set_title(session_id, title, TitleSource.CUSTOM)
671 except SessionNotFoundError as exc:
672 return _error(str(exc))
673 return {"id": session_id, "title": title}
676@_tool_if(agent_sessions_enabled)
677def session_delete(session_id: str) -> dict[str, Any]:
678 """Delete an agent session."""
679 if not agent_sessions_enabled():
680 return _error(AGENT_SESSIONS_DISABLED_HINT)
681 try:
682 _require_agent_session(session_id)
683 get_services().session_store.delete(session_id)
684 except SessionNotFoundError as exc:
685 return _error(str(exc))
686 return {"id": session_id, "deleted": True}
689@_tool
690def export_dataset(output: str, fmt: str = "", source: str = "") -> dict[str, Any]:
691 """Write the per-page {source, page, text} dataset to a file (no vectors).
693 ``fmt`` is parquet/jsonl (empty infers from the suffix); ``source`` limits to one file.
694 """
695 from lilbee.app.dataset import DatasetError, export_to_path
697 try:
698 summary = export_to_path(Path(output), fmt, source or None, cancel=_caller_cancelled())
699 except DatasetError as exc:
700 return _error(str(exc))
701 return summary.model_dump()
704@_tool
705async def import_dataset(dataset: str, fmt: str = "", ctx: Context | None = None) -> dict[str, Any]:
706 """Import a per-page text dataset, re-embedding under the current model.
708 Replaces existing copies; imported sources are detached so sync won't delete them.
709 """
710 from lilbee.app.dataset import DatasetError, import_from_path
711 from lilbee.runtime.progress import EmbedEvent, EventType, ProgressEvent
713 loop = asyncio.get_running_loop()
715 with _cancel_token() as cancel:
717 def on_progress(event_type: EventType, data: ProgressEvent) -> None:
718 # Raise rather than return: the embed work runs off the loop, so a
719 # quiet return would re-embed the whole dataset for a caller that
720 # has gone. Checked before the isinstance filter so the stop lands
721 # on every event, not only the ones that map to a percent.
722 if cancel.is_set():
723 raise TaskCancelledError
724 # EMBED events carry chunk/total_chunks; other event types don't map to a percent.
725 if ctx is None or not isinstance(data, EmbedEvent):
726 return
727 future = asyncio.run_coroutine_threadsafe(
728 ctx.report_progress(
729 progress=float(data.chunk), total=float(data.total_chunks), message=data.file
730 ),
731 loop,
732 )
733 future.add_done_callback(_log_progress_failure)
735 try:
736 summary = await import_from_path(Path(dataset), fmt, on_progress=on_progress)
737 except DatasetError as exc:
738 return _error(str(exc))
739 return summary.model_dump()
742@_tool
743def reset(confirm: bool = False) -> dict[str, Any]:
744 """Factory reset: delete all documents and indexed data. Requires ``confirm=true``."""
745 if _transport.http_mounted:
746 return _error(
747 "reset is unavailable on the HTTP server: it would wipe the shared index for "
748 "every connected client. Run it from the CLI or the stdio MCP server."
749 )
750 if not confirm:
751 return _error("pass confirm=true to confirm deletion")
752 from lilbee.app.reset import perform_reset
754 result = perform_reset().model_dump()
755 # Reopen LanceDB against the empty data dir; keep providers loaded.
756 reset_store()
757 return result
760@_wiki_tool
761def wiki_lint(wiki_source: str = "") -> dict[str, Any]:
762 """Lint wiki pages; empty ``wiki_source`` lints all."""
763 from lilbee.wiki.lint import LintReport, lint_all, lint_wiki_page
765 store = get_services().store
766 report = (
767 LintReport(issues=lint_wiki_page(wiki_source, store))
768 if wiki_source
769 else lint_all(store, cancel=_caller_cancelled())
770 )
771 return {
772 "command": "wiki_lint",
773 "issues": [i.to_dict() for i in report.issues],
774 "total": len(report.issues),
775 "errors": report.error_count,
776 "warnings": report.warning_count,
777 }
780@_wiki_tool
781def wiki_citations(wiki_source: str = "", source: str = "") -> dict[str, Any]:
782 """List a wiki page's citations, or with ``source``, the pages citing that document.
784 Pass exactly one of ``wiki_source`` (forward) or ``source`` (reverse).
785 """
786 if bool(wiki_source) == bool(source):
787 return _error("pass either wiki_source or source, not both")
788 store = get_services().store
789 if source:
790 records = store.get_citations_for_source(source)
791 return {
792 "command": "wiki_citations",
793 "source": source,
794 "citations": [dict(r) for r in records],
795 "total": len(records),
796 }
797 records = store.get_citations_for_wiki(wiki_source)
798 return {
799 "command": "wiki_citations",
800 "wiki_source": wiki_source,
801 "citations": [dict(r) for r in records],
802 "total": len(records),
803 }
806@_tool
807def wiki_status() -> dict[str, Any]:
808 """Show wiki layer status: page counts, recent lint issues.
810 Registered even when the wiki is disabled, like the HTTP status route, so a
811 caller can read the disabled state instead of finding no tool at all.
812 """
813 from lilbee.wiki.lint import lint_all
815 wiki_root = cfg.data_root / cfg.wiki_dir
816 if not cfg.wiki or not wiki_root.exists():
817 # Same keys as the enabled arm so a client can key on lint_errors in
818 # either state; a disabled wiki reports zeros rather than being linted.
819 return {
820 "wiki_enabled": cfg.wiki,
821 WikiSubdir.SUMMARIES: 0,
822 WikiSubdir.DRAFTS: 0,
823 "pages": 0,
824 "lint_errors": 0,
825 "lint_warnings": 0,
826 }
828 summaries_dir = wiki_root / WikiSubdir.SUMMARIES
829 drafts_dir = wiki_root / WikiSubdir.DRAFTS
830 summaries = list(summaries_dir.rglob("*.md")) if summaries_dir.exists() else []
831 drafts = list(drafts_dir.rglob("*.md")) if drafts_dir.exists() else []
833 # Read-only status: lint for counts without appending to the audit log.
834 report = lint_all(get_services().store, record_log=False, cancel=_caller_cancelled())
835 return {
836 "wiki_enabled": cfg.wiki,
837 WikiSubdir.SUMMARIES: len(summaries),
838 WikiSubdir.DRAFTS: len(drafts),
839 "pages": total_wiki_pages(wiki_root),
840 "lint_errors": report.error_count,
841 "lint_warnings": report.warning_count,
842 }
845@_wiki_tool
846def wiki_list() -> dict[str, Any]:
847 """List wiki pages with metadata."""
848 from dataclasses import asdict
850 from lilbee.wiki.browse import list_pages
852 wiki_root = cfg.data_root / cfg.wiki_dir
853 pages = list_pages(wiki_root)
854 return {
855 "command": "wiki_list",
856 "pages": [asdict(p) for p in pages],
857 "total": len(pages),
858 }
861@_wiki_tool
862def wiki_read(slug: str) -> dict[str, Any]:
863 """Read a wiki page's content + frontmatter by slug."""
864 from dataclasses import asdict
866 from lilbee.wiki.browse import read_page
868 wiki_root = cfg.data_root / cfg.wiki_dir
869 result = read_page(wiki_root, slug)
870 if result is None:
871 return _error(f"wiki page not found: {slug}")
872 return {"command": "wiki_read", **asdict(result)}
875@_wiki_tool
876def wiki_build(dry_run: bool = False) -> dict[str, Any]:
877 """Build the concept and entity wiki across all ingested sources. Blocks until done.
879 ``dry_run=True`` returns the NER entity candidates a build would cover and
880 makes no LLM call.
881 """
882 from lilbee.wiki import run_full_build
883 from lilbee.wiki.generation import DRY_RUN_CONCEPT_NOTE, preview_build_entities
885 if dry_run:
886 rows = preview_build_entities(cfg)
887 return {
888 "command": "wiki_build",
889 "dry_run": True,
890 "entities": rows,
891 "count": len(rows),
892 "note": DRY_RUN_CONCEPT_NOTE,
893 }
894 return {"command": "wiki_build", **run_full_build(cfg, cancel=_caller_cancelled())}
897@_wiki_tool
898def wiki_update() -> dict[str, Any]:
899 """Refresh the concept and entity wiki after an ingest. A full rebuild; blocks until done."""
900 from lilbee.wiki import run_full_build
902 return {"command": "wiki_update", **run_full_build(cfg, cancel=_caller_cancelled())}
905@_wiki_tool
906def wiki_synthesize() -> dict[str, Any]:
907 """Generate synthesis pages for concept clusters with three or more sources."""
908 from lilbee.wiki import run_full_synthesize
910 return {"command": "wiki_synthesize", **run_full_synthesize(cfg, cancel=_caller_cancelled())}
913@_wiki_tool
914def wiki_prune() -> dict[str, Any]:
915 """Prune stale and orphaned wiki pages."""
916 from lilbee.wiki.prune import prune_wiki
918 report = prune_wiki(get_services().store, cancel=_caller_cancelled())
919 return {
920 "command": "wiki_prune",
921 "records": [r.to_dict() for r in report.records],
922 "archived": report.archived_count,
923 "flagged": report.flagged_count,
924 "reconciled": report.reconciled_count,
925 }
928@_wiki_tool
929def wiki_index() -> dict[str, Any]:
930 """Rebuild the browse index of pages the corpus could have. No LLM call."""
931 from lilbee.wiki.stubs import refresh_stub_index
933 stubs = refresh_stub_index(get_services().store)
934 return {"command": "wiki_index", "entries": len(stubs)}
937@_wiki_tool
938def wiki_generate(slug: str) -> dict[str, Any]:
939 """Generate one indexed wiki page. Costs a single LLM call and is GPU-heavy."""
940 from lilbee.wiki.browse import page_slug
941 from lilbee.wiki.lazy import UnknownStubError, generate_stub_page
943 try:
944 path = generate_stub_page(slug, get_services().store)
945 except UnknownStubError as exc:
946 return _error(str(exc))
947 if path is None:
948 return _error(f"index entry for {slug} is stale; its sources are gone")
949 # The read surfaces address pages by section, so answer with that slug.
950 read_slug = page_slug(path, cfg.data_root / cfg.wiki_dir)
951 return {"command": "wiki_generate", "slug": read_slug, "path": path.as_posix()}
954@_tool
955def wiki_wipe(confirm: bool = False) -> dict[str, Any]:
956 """Delete every generated wiki page and its indexed rows. Pass ``confirm=true``.
958 Registered even with the wiki disabled, because turning the setting off
959 leaves the pages generated earlier in place.
960 """
961 if not confirm:
962 return _error("pass confirm=true to delete the wiki; this cannot be undone")
963 from lilbee.wiki.wipe import wipe_wiki
965 report = wipe_wiki(get_services().store)
966 if not report.rows_deleted:
967 return _error(report.summary())
968 return {
969 "command": "wiki_wipe",
970 "pages_removed": report.pages_removed,
971 "sources_cleared": report.sources_cleared,
972 }
975def _setting_info_to_dict(info: SettingInfo) -> dict[str, Any]:
976 """Render a SettingInfo as a JSON-safe dict for the MCP wire format."""
977 return {
978 "key": info.key,
979 "value": _json_safe(info.value),
980 "default": _json_safe(info.default),
981 "type": info.type,
982 "nullable": info.nullable,
983 "group": info.group.value,
984 "help": info.help_text,
985 "choices": list(info.choices) if info.choices else None,
986 "reindex_required": info.reindex_required,
987 }
990def _json_safe(value: Any) -> Any:
991 """Coerce Path / frozenset / tuple to JSON-friendly primitives."""
992 if isinstance(value, str | int | float | bool | list | type(None)):
993 return value
994 return str(value)
997@_tool
998def settings_list(group: str = "") -> dict[str, Any]:
999 """List writable lilbee settings (each with value, default, type, help, choices).
1001 ``group`` filters by group name (case-insensitive); empty returns all.
1002 """
1004 try:
1005 infos = list_settings(group or None)
1006 except ValueError as exc:
1007 return _error(str(exc))
1008 return {
1009 "command": "settings_list",
1010 "settings": [_setting_info_to_dict(info) for info in infos],
1011 "total": len(infos),
1012 }
1015@_tool
1016def settings_get(key: str) -> dict[str, Any]:
1017 """Get a single setting's current value + metadata."""
1019 try:
1020 info = get_setting(key)
1021 except KeyError as exc:
1022 return _error(str(exc))
1023 return {"command": "settings_get", "setting": _setting_info_to_dict(info)}
1026@_tool
1027def settings_set(updates: dict[str, Any]) -> dict[str, Any]:
1028 """Atomically update writable settings; rolls back on validation error.
1029 Persists to config.toml; returns ``{updated, reindex_required}``."""
1030 if _transport.http_mounted and requires_services_reset(updates):
1031 return _error(provider_reset_refused_message("Switching"))
1032 try:
1033 result = apply_settings_update(updates)
1034 except (ValueError, TypeError) as exc:
1035 return _error(str(exc))
1036 return {
1037 "command": "settings_set",
1038 "updated": result.updated,
1039 "reindex_required": result.reindex_required,
1040 }
1043@_tool
1044def settings_reset(keys: list[str]) -> dict[str, Any]:
1045 """Reset writable settings to their built-in defaults."""
1046 if _transport.http_mounted and requires_services_reset(dict.fromkeys(keys)):
1047 return _error(provider_reset_refused_message("Resetting"))
1048 try:
1049 result = reset_settings(keys)
1050 except (ValueError, TypeError) as exc:
1051 return _error(str(exc))
1052 return {
1053 "command": "settings_reset",
1054 "updated": result.updated,
1055 "reindex_required": result.reindex_required,
1056 }
1059@_tool
1060def model_list(source: str = "", task: str = "") -> dict[str, Any]:
1061 """List installed models. ``source`` is ``native`` / ``remote``; ``task`` filters by role."""
1062 from lilbee.app.models import list_models_data
1063 from lilbee.catalog.types import ModelTask
1065 try:
1066 src = ModelSource.parse(source)
1067 except ValueError as exc:
1068 return _error(str(exc))
1069 try:
1070 parsed_task = ModelTask(task) if task else None
1071 except ValueError as exc:
1072 return _error(str(exc))
1073 return list_models_data(source=src, task=parsed_task).model_dump()
1076@_tool
1077def catalog_browse(
1078 task: str = "",
1079 search: str = "",
1080 size: str = "",
1081 installed: bool | None = None,
1082 featured: bool | None = None,
1083 sort: str = "featured",
1084 limit: int = 20,
1085 offset: int = 0,
1086) -> dict[str, Any]:
1087 """Browse the lilbee model catalog. ``task``: chat/embedding/vision/rerank.
1088 ``size``: small/medium/large/huge, by parameter count.
1089 ``sort``: featured/downloads/name/size_asc/size_desc."""
1090 from lilbee.catalog.query import get_catalog
1091 from lilbee.catalog.types import CatalogSize, CatalogSort, ModelTask
1093 try:
1094 parsed_task = ModelTask(task) if task else None
1095 parsed_size = CatalogSize(size) if size else None
1096 parsed_sort = CatalogSort(sort)
1097 except ValueError as exc:
1098 return _error(str(exc))
1099 try:
1100 result = get_catalog(
1101 task=parsed_task,
1102 search=search,
1103 size=parsed_size,
1104 installed=installed,
1105 featured=featured,
1106 sort=parsed_sort,
1107 limit=limit,
1108 offset=offset,
1109 model_manager=get_services().model_manager,
1110 )
1111 except ValueError as exc:
1112 return _error(str(exc))
1113 return {
1114 "command": "catalog_browse",
1115 "total": result.total,
1116 "limit": result.limit,
1117 "offset": result.offset,
1118 "has_more": result.has_more,
1119 "models": [
1120 {
1121 "ref": m.hf_repo,
1122 "display_name": m.display_name,
1123 "task": m.task.value,
1124 "size_gb": m.size_gb,
1125 "min_ram_gb": m.min_ram_gb,
1126 "downloads": m.downloads,
1127 "featured": m.featured,
1128 "description": m.description,
1129 "architecture": m.architecture,
1130 "compat": m.compat.value,
1131 }
1132 for m in result.models
1133 ],
1134 }
1137@_tool
1138def model_show(model: str) -> dict[str, Any]:
1139 """Show catalog and installed metadata for a model ref."""
1140 from lilbee.app.models import show_model_data
1141 from lilbee.modelhub.model_manager import ModelNotFoundError
1143 try:
1144 return show_model_data(model).model_dump()
1145 except ModelNotFoundError as exc:
1146 return _error(str(exc))
1149def _log_progress_failure(future: concurrent.futures.Future[None]) -> None:
1150 """Log report_progress failures without raising.
1152 Progress notifications are best-effort: a failure should not abort
1153 an in-flight pull.
1154 """
1155 try:
1156 future.result()
1157 except Exception:
1158 log.warning("MCP report_progress failed", exc_info=True)
1161@_tool
1162async def model_pull(
1163 model: str,
1164 source: str = ModelSource.NATIVE.value,
1165 allow_unsupported: bool = False,
1166 ctx: Context | None = None,
1167) -> dict[str, Any]:
1168 """Download a model and stream progress.
1170 ``source`` is ``native`` (GGUF) or ``remote`` (SDK).
1171 ``allow_unsupported`` overrides the supported-architecture refusal.
1172 """
1173 from lilbee.app.models import pull_model_data
1174 from lilbee.catalog import DownloadProgress
1175 from lilbee.catalog.compat import SUPPORTED_ARCHS, UnsupportedArchError
1177 try:
1178 src = ModelSource.parse(source) or ModelSource.NATIVE
1179 except ValueError as exc:
1180 return _error(str(exc))
1182 loop = asyncio.get_running_loop()
1184 with _cancel_token() as cancel:
1186 def on_update(p: DownloadProgress) -> None:
1187 # Raise rather than return: the download runs on a thread asyncio
1188 # cannot interrupt, so returning would leave a multi-GB pull going
1189 # for a caller that has gone. Same idiom the HTTP pull uses.
1190 if cancel.is_set():
1191 raise TaskCancelledError
1192 if ctx is None:
1193 return
1194 future = asyncio.run_coroutine_threadsafe(
1195 ctx.report_progress(progress=float(p.percent), total=100.0, message=p.detail),
1196 loop,
1197 )
1198 future.add_done_callback(_log_progress_failure)
1200 try:
1201 result = await asyncio.to_thread(
1202 pull_model_data,
1203 model,
1204 src,
1205 on_update=on_update,
1206 allow_unsupported=allow_unsupported,
1207 )
1208 except UnsupportedArchError as exc:
1209 return {
1210 "ok": False,
1211 "command": "model_pull",
1212 "error": {
1213 "code": "unsupported_arch",
1214 "arch": exc.architecture,
1215 "ref": exc.ref,
1216 "supported_examples": sorted(SUPPORTED_ARCHS)[:5],
1217 "total_supported": len(SUPPORTED_ARCHS),
1218 },
1219 }
1220 except (RuntimeError, PermissionError) as exc:
1221 return _error(str(exc))
1222 return result.model_dump()
1225@_tool
1226def model_rm(model: str, source: str = "") -> dict[str, Any]:
1227 """Remove an installed model. Only native GGUF models lilbee downloaded;
1228 Ollama/LM Studio are read-only."""
1229 from lilbee.app.models import remove_model_data
1231 try:
1232 src = ModelSource.parse(source)
1233 return remove_model_data(model, source=src).model_dump()
1234 except ValueError as exc:
1235 return _error(str(exc))
1238@_wiki_tool
1239def wiki_drafts_list() -> dict[str, Any]:
1240 """List pending wiki drafts. Read-only: promotion is reserved for the human
1241 surfaces (CLI, TUI, and the authenticated HTTP API)."""
1242 from lilbee.wiki.drafts import list_drafts
1244 wiki_root = cfg.data_root / cfg.wiki_dir
1245 drafts = list_drafts(wiki_root)
1246 return {
1247 "command": "wiki_drafts_list",
1248 "drafts": [d.to_dict() for d in drafts],
1249 "total": len(drafts),
1250 }
1253@_wiki_tool
1254def wiki_drafts_diff(slug: str) -> dict[str, Any]:
1255 """Unified diff of a draft against its published counterpart."""
1256 from lilbee.core.security import PathTraversalError
1257 from lilbee.wiki.drafts import diff_draft
1259 wiki_root = cfg.data_root / cfg.wiki_dir
1260 try:
1261 diff = diff_draft(slug, wiki_root)
1262 except FileNotFoundError as exc:
1263 return _error(str(exc))
1264 except PathTraversalError:
1265 return _error(INVALID_DRAFT_SLUG_ERROR)
1266 return {"command": "wiki_drafts_diff", "slug": slug, "diff": diff}
1269def _collapse_nullable_anyof(prop: dict[str, Any]) -> None:
1270 """Collapse ``anyOf: [{type: X}, {type: null}]`` to ``{type: X}`` in place.
1272 Pydantic emits ``T | None`` parameters as a two-arm anyOf with a null
1273 branch. The null branch carries no information the model needs to pick
1274 or shape its call, but it costs tokens at every dispatch. Drop it.
1275 """
1276 arms = prop.get("anyOf")
1277 if not isinstance(arms, list):
1278 return
1279 non_null = [a for a in arms if isinstance(a, dict) and a.get("type") != "null"]
1280 if len(non_null) == 1 and len(non_null) < len(arms):
1281 prop.pop("anyOf", None)
1282 for key, value in non_null[0].items():
1283 prop.setdefault(key, value)
1286def _strip_property_noise(prop: dict[str, Any]) -> None:
1287 """Drop tokens that don't change the model's behavior."""
1288 prop.pop("title", None)
1289 prop.pop("default", None)
1290 _collapse_nullable_anyof(prop)
1291 if prop.get("additionalProperties") is True:
1292 prop.pop("additionalProperties", None)
1295def _flatten_tool_description(text: str) -> str:
1296 """Flatten a triple-quoted tool docstring for the tools wire.
1298 The summary line carries no indent while continuation lines are indented to
1299 the source, so ``textwrap.dedent`` alone is a no-op (the common prefix is the
1300 empty string) and leaves source indentation on every body line -- including
1301 deeper-indented Args lines. Strip each line so the model sees flat text;
1302 blank lines are kept so paragraph breaks survive.
1303 """
1304 return "\n".join(line.strip() for line in text.strip().splitlines())
1307def _strip_schema(schema: dict[str, Any]) -> dict[str, Any]:
1308 """Trim auto-generated noise from a tool's input schema, on a copy.
1310 Drops:
1311 - SDK/Pydantic ``title`` keys (per-schema + per-property). Tools the
1312 model picks by name don't need a separate display title.
1313 - ``default`` values on properties: clients send what they want and
1314 omitted fields fall back server-side.
1315 - ``additionalProperties: true`` on dict params: Pydantic emits it for
1316 every ``dict[str, Any]`` but it's the JSON Schema default behavior.
1317 - The ``null`` arm of ``anyOf: [{type: X}, {type: null}]`` unions for
1318 ``T | None`` defaults; the null branch is implicit.
1320 A roughly 25-35% reduction in the serialized tools payload, which matters
1321 most for small-context (16K) chat models where the tools surface was
1322 previously eating ~60% of the budget.
1323 """
1324 schema = deepcopy(schema)
1325 schema.pop("title", None)
1326 properties = schema.get("properties")
1327 if isinstance(properties, dict):
1328 for prop in properties.values():
1329 if isinstance(prop, dict):
1330 _strip_property_noise(prop)
1331 return schema
1334_NO_WIKI_SCOPE_HINT = ' No wiki layer here: use scope "raw" or "both".'
1337class LilbeeMCP(MCPServer):
1338 """MCP server that trims its tools wire and keeps it current with config."""
1340 async def list_tools(self) -> list[MCPTool]:
1341 """The registered tools with schema noise stripped and flat descriptions.
1343 The transforms run on the wire representation per request, never on the
1344 stored registrations, so they cannot drift out of sync with config. The
1345 ``search`` description advertises only the scopes this corpus has: when
1346 wiki generation is off, a model that guesses ``scope="wiki"`` gets a
1347 silent fallback to the full pool, so raw/both only.
1348 """
1349 tools = await super().list_tools()
1350 for tool in tools:
1351 tool.input_schema = _strip_schema(tool.input_schema)
1352 if isinstance(tool.description, str):
1353 description = _flatten_tool_description(tool.description)
1354 if tool.name == "search" and not cfg.wiki:
1355 description += _NO_WIKI_SCOPE_HINT
1356 tool.description = description
1357 return tools
1360def _client_name(ctx: Context | None) -> str:
1361 """The MCP client's self-reported name from the initialize handshake, or empty."""
1362 if ctx is None:
1363 return ""
1364 params = ctx.session.client_params
1365 return params.clientInfo.name if params is not None else ""
1368def _slug(value: str) -> str:
1369 """Lowercase, hyphenated id fragment; falls back to ``generic`` when empty."""
1370 slug = re.sub(r"[^a-z0-9]+", "-", value.lower()).strip("-")
1371 return slug or "generic"
1374# Per-connection fallback ids for agents that report no identity. Keyed by the
1375# live MCP session so each connection gets a distinct, stable namespace instead
1376# of every unidentified agent colliding on a shared one. WeakKeyDictionary drops
1377# entries once the session is collected, so this does not grow unbounded. The
1378# lock guards the get-or-create because sync tool handlers run on the offload
1379# threadpool, so concurrent connections (and weakref-removal callbacks) can
1380# touch the mapping from different threads.
1381_ANON_OWNER_IDS: WeakKeyDictionary[object, str] = WeakKeyDictionary()
1382_ANON_OWNER_LOCK = threading.Lock()
1385def _anon_owner_id(ctx: Context | None) -> str:
1386 """A stable per-connection id for an agent that reported no identity.
1388 Without this, two unidentified agents would both slug to ``generic`` and
1389 share a memory namespace; keying on the session keeps them isolated.
1390 """
1391 if ctx is None:
1392 return "anonymous"
1393 session = ctx.session
1394 with _ANON_OWNER_LOCK:
1395 existing = _ANON_OWNER_IDS.get(session)
1396 if existing is None:
1397 existing = f"anon-{uuid.uuid4().hex[:12]}"
1398 _ANON_OWNER_IDS[session] = existing
1399 return existing
1402def _derive_owner(agent_id: str, ctx: Context | None) -> str:
1403 """Resolve the calling agent's stable owner namespace.
1405 Precedence: explicit ``agent_id`` argument, then the ``LILBEE_AGENT_ID`` env var
1406 (pinned in the client's MCP config), then the MCP client name, then a stable
1407 per-connection fallback so unidentified agents never share a namespace.
1408 """
1409 explicit = agent_id or os.environ.get("LILBEE_AGENT_ID", "")
1410 resolved = explicit or _client_name(ctx)
1411 if resolved:
1412 return agent_owner(_slug(resolved))
1413 return agent_owner(_slug(_anon_owner_id(ctx)))
1416@_tool_if(memory_enabled)
1417def memory_remember(
1418 text: str,
1419 kind: MemoryKind = MemoryKind.FACT,
1420 shared: bool = False,
1421 agent_id: str = "",
1422 ctx: Context | None = None,
1423) -> dict[str, Any]:
1424 """Store a durable memory. ``kind``: "fact" (similarity-recalled) or "preference" (always on).
1425 ``shared`` exposes it to the human's TUI/CLI."""
1426 if not memory_enabled():
1427 return _error(MEMORY_DISABLED_HINT)
1428 owner = _derive_owner(agent_id, ctx)
1429 memory_id = remember(text, owner=owner, kind=kind, source=MemorySource.AGENT, shared=shared)
1430 return {"ok": True, "id": memory_id, "owner": owner}
1433@_tool_if(memory_enabled)
1434def memory_recall(
1435 query: str, limit: int = 0, agent_id: str = "", ctx: Context | None = None
1436) -> dict[str, Any]:
1437 """Recall this agent's memories (plus any the human shared) relevant to *query*."""
1438 if not memory_enabled():
1439 return _error(MEMORY_DISABLED_HINT)
1440 owner = _derive_owner(agent_id, ctx)
1441 memories = recall(query, owner, top_k=limit if limit > 0 else None)
1442 return {
1443 "memories": [
1444 {"id": m.id, "text": m.text, "kind": m.kind.value, "owner": m.owner} for m in memories
1445 ]
1446 }
1449@_tool_if(memory_enabled)
1450def memory_list(agent_id: str = "", ctx: Context | None = None) -> dict[str, Any]:
1451 """List every memory in this agent's namespace (any kind, newest first)."""
1452 if not memory_enabled():
1453 return _error(MEMORY_DISABLED_HINT)
1454 owner = _derive_owner(agent_id, ctx)
1455 memories = list_memories(owner)
1456 return {
1457 "memories": [
1458 {"id": m.id, "text": m.text, "kind": m.kind.value, "shared": m.shared} for m in memories
1459 ]
1460 }
1463@_tool_if(memory_enabled)
1464def memory_forget(memory_id: str, agent_id: str = "", ctx: Context | None = None) -> dict[str, Any]:
1465 """Delete one of this agent's own memories by id (agent_id scopes the namespace)."""
1466 if not memory_enabled():
1467 return _error(MEMORY_DISABLED_HINT)
1468 owner = _derive_owner(agent_id, ctx)
1469 if not forget(memory_id, owner=owner):
1470 return _error(f"No memory '{memory_id}' in this agent's namespace.")
1471 return {"ok": True, "id": memory_id}
1474def _placement_dict(view: PlacementView) -> dict[str, Any]:
1475 from lilbee.server.models import PlacementResponse
1477 return PlacementResponse.from_view(view).model_dump(mode="json")
1480def _placement_guard(serialize: Callable[[], dict[str, Any]]) -> dict[str, Any]:
1481 """Run a placement query and serialize it, returning a structured error on failure."""
1482 from lilbee.providers.base import ProviderError
1483 from lilbee.providers.fleet.placement_spec import PlacementError
1485 try:
1486 return serialize()
1487 except (PlacementError, ProviderError) as exc:
1488 return _error(str(exc))
1491def _placement_result(action: Callable[[], PlacementView]) -> dict[str, Any]:
1492 """Run a placement action and serialize its view, returning a structured error on failure."""
1493 return _placement_guard(lambda: _placement_dict(action()))
1496def _parse_spec(spec: dict[str, Any] | None) -> PlacementSpec | None:
1497 from lilbee.providers.fleet.placement_spec import PlacementSpec
1499 return PlacementSpec.from_json(json.dumps(spec)) if spec else None
1502@_tool_named("get_gpus")
1503def get_gpus_tool() -> dict[str, Any]:
1504 """List detected GPUs with free/total VRAM (the placement HTTP /api/gpus equivalent)."""
1506 def _body() -> dict[str, Any]:
1507 from lilbee.cli.tui import messages as msg
1508 from lilbee.providers.fleet.gpu_stats import probe_intel_util_hint
1510 view = get_placement()
1511 hint = probe_intel_util_hint(view.gpus)
1512 return {
1513 "gpus": _placement_dict(view)["gpus"],
1514 "notice": msg.intel_util_hint_text(hint) if hint else None,
1515 }
1517 return _placement_guard(_body)
1520@_tool_named("get_placement")
1521def get_placement_tool() -> dict[str, Any]:
1522 """Show the current effective multi-GPU model placement."""
1523 return _placement_result(get_placement)
1526@_tool_named("preview_placement")
1527def preview_placement_tool(spec: dict[str, Any] | None = None) -> dict[str, Any]:
1528 """Preview what a placement spec (or auto, when omitted) would place. No changes made."""
1529 return _placement_result(lambda: preview_placement(_parse_spec(spec)))
1532@_tool_named("set_placement")
1533def set_placement_tool(spec: dict[str, Any]) -> dict[str, Any]:
1534 """Set and apply a manual multi-GPU placement spec (persists to config).
1536 The spec maps a role ("chat"/"embed"/"rerank"/"vision") to a placement, e.g.
1537 ``{"chat": {"devices": [0, 1], "tensor_split": [1, 1]}}``. ``devices`` is the
1538 GPU indices (get_gpus lists them); ``tensor_split`` is optional per-device
1539 weights (omit for an even split). Omit a role to leave it auto-placed.
1540 """
1541 from lilbee.providers.fleet.placement_spec import PlacementSpec
1543 # set_placement restarts the shared fleet's moved roles: gate it on the
1544 # shared HTTP transport exactly like the REST PUT/DELETE placement routes.
1545 if _transport.http_mounted and not cfg.allow_http_placement:
1546 return _error(placement_refused_message())
1547 # Always build a spec (even {}) so an empty/invalid one is rejected, not cleared.
1548 return _placement_result(lambda: set_placement(PlacementSpec.from_json(json.dumps(spec))))
1551@_tool_named("clear_placement")
1552def clear_placement_tool() -> dict[str, Any]:
1553 """Clear the manual placement and return to automatic placement."""
1554 if _transport.http_mounted and not cfg.allow_http_placement:
1555 return _error(placement_refused_message())
1556 return _placement_result(lambda: set_placement(None))
1559def build_mcp_server() -> LilbeeMCP:
1560 """Build an MCP server carrying every tool registered in this module.
1562 Each transport builds its own instance: the SDK caches one
1563 ``StreamableHTTPSessionManager`` per server and its ``run()`` is single-use,
1564 so a shared server cannot back two apps in one process. Gates registered
1565 via ``_tool_if`` are evaluated here, against current config.
1566 """
1567 server = LilbeeMCP("lilbee", instructions=_INSTRUCTIONS)
1568 for fn, name, gate in _REGISTRATIONS:
1569 if gate is None or gate():
1570 server.add_tool(fn, name=name)
1571 return server
1574_PARENT_DEATH_CLEANUP_S = 5.0
1577def _exit_on_parent_death() -> None:
1578 """Release engine membership best-effort, then hard-exit promptly.
1580 ``os._exit`` skips atexit, so an explicit release stops this process's engine
1581 when it was the last user and keeps the machine clean. But this watchdog's one
1582 contract is to exit promptly on parent death, so the release runs on a daemon
1583 thread joined with a short deadline: a peer holding an engine build lock can
1584 never keep the orphaned process alive with its models resident. The kernel
1585 releases this process's user lock on exit regardless, so a skipped release only
1586 defers the engine stop to the peers' reap and the idle TTL.
1587 """
1588 cleanup = threading.Thread(target=reset_services, name="parent-death-cleanup", daemon=True)
1589 cleanup.start()
1590 cleanup.join(timeout=_PARENT_DEATH_CLEANUP_S)
1591 os._exit(0)
1594def main() -> None:
1595 """Entry point for the stdio MCP server."""
1596 # Preload so the first tool call doesn't pay the cold-start cost
1597 # of provider/embedder/store init. Failures (missing model, bad
1598 # config) still surface on the first tool call rather than crashing
1599 # the server before it attaches to stdio.
1600 try:
1601 get_services()
1602 except Exception:
1603 log.debug("MCP pre-warm failed; services will init on first call", exc_info=True)
1605 from lilbee.parent_monitor import parse_parent_pid, watch_parent_thread
1607 parent_pid = parse_parent_pid()
1608 if parent_pid is not None:
1609 watch_parent_thread(parent_pid, _exit_on_parent_death)
1611 build_mcp_server().run()