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