Coverage for src/lilbee/mcp_server.py: 100%
959 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +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, dataclass
20from enum import StrEnum
21from pathlib import Path
22from typing import TYPE_CHECKING, Annotated, Any, TypeVar, cast
23from weakref import WeakKeyDictionary
25import anyio
26from mcp.server.mcpserver import Context, MCPServer
27from mcp.types import Tool as MCPTool
28from pydantic import Field
30from lilbee.app import profiles
31from lilbee.app.analyze import (
32 AnalyzeRequest,
33 hide_tip,
34 run_analysis,
35 server_directory,
36 tip_state,
37)
38from lilbee.app.config_follow import ConfigFileFollower
39from lilbee.app.memory import (
40 MEMORY_DISABLED_HINT,
41 forget,
42 list_memories,
43 memory_enabled,
44 recall,
45 remember,
46)
47from lilbee.app.placement import (
48 PlacementView,
49 get_placement,
50 placement_refused_message,
51 preview_placement,
52 set_placement,
53)
54from lilbee.app.search import clean_results
55from lilbee.app.services import (
56 exit_on_signals_in_loop,
57 get_services,
58 reset_services,
59 reset_store,
60 set_config_refresh,
61)
62from lilbee.app.session_sources import (
63 LocatorError,
64 SourceRefInput,
65 message_payload,
66 posted_message,
67)
68from lilbee.app.settings import (
69 SettingInfo,
70 apply_settings_update,
71 config_write_failure_message,
72 get_setting,
73 list_settings,
74 provider_reset_refused_message,
75 requires_services_reset,
76 reset_settings,
77)
78from lilbee.catalog.types import ModelSource
79from lilbee.core.config import cfg, validate_ocr_timeout
80from lilbee.core.config.enums import CrawlRenderMode, OcrMode
81from lilbee.core.profile_files import ProfileFolder, ProfileStore, ProfileValidation
82from lilbee.core.settings import overlay_persisted_settings
83from lilbee.core.system import LOCAL_ROOT_DIRNAME, canonical_data_root
84from lilbee.crawler import (
85 CRAWLER_MISSING_MESSAGE,
86 crawler_available,
87 is_url,
88 require_valid_crawl_url,
89)
90from lilbee.crawler.bootstrap import CrawlEngineRefusedError
91from lilbee.crawler.models import CrawlLogin, crawl_limit_error
92from lilbee.crawler.stored_login import STORED_LOGIN_NOT_SENT, crawl_login, stored_login_withheld
93from lilbee.crawler.task import get_task, start_crawl
94from lilbee.crawler.url_filter import without_userinfo
95from lilbee.data.store import (
96 DOCUMENTS_PAGE_DEFAULT,
97 DOCUMENTS_PAGE_MAX,
98 EmbeddingModelMismatchError,
99 MemoryKind,
100 MemorySource,
101 SearchScope,
102 agent_owner,
103 scope_to_chunk_type,
104)
105from lilbee.runtime.cancellation import TaskCancelledError
106from lilbee.runtime.hardware import FitLevel, available_memory_for_fit, make_fit_filter
107from lilbee.runtime.lock import ResetRefusedError
108from lilbee.runtime.progress import (
109 UPGRADE_HEADLINE,
110 CrawlPageFailedEvent,
111 EmbedEvent,
112 EventType,
113 ProgressEvent,
114)
115from lilbee.sessions import (
116 AGENT_SESSIONS_DISABLED_HINT,
117 MessageRole,
118 Session,
119 SessionForkRangeError,
120 SessionNotFoundError,
121 SessionOrigin,
122 SessionOwnershipError,
123 StoredSourceRef,
124 TitleSource,
125 agent_sessions_enabled,
126)
127from lilbee.wiki.shared import (
128 INVALID_DRAFT_SLUG_ERROR,
129 WIKI_DISABLED_ERROR,
130 WikiSubdir,
131 total_wiki_pages,
132)
134if TYPE_CHECKING:
135 from lilbee.providers.fleet.placement_spec import PlacementSpec
137log = logging.getLogger(__name__)
138# What a sync of this server is called in the notice of a sync that waits for it.
139_SYNC_STARTER = "lilbee mcp"
141_INSTRUCTIONS = (
142 "Local search engine over the user's files, code, and crawled pages. "
143 "For any question about the user's own documents or codebase -- a lookup, a "
144 "find-in-docs, 'where is X', 'how does Y work here' -- call lilbee_search first "
145 "and answer from its cited chunks. Prefer it over web-fetch or file-read tools: "
146 "those cannot see the indexed corpus."
147)
150class _TransportState:
151 """Process-level MCP transport facts.
153 ``http_mounted`` is True only when MCP is exposed over the shared
154 streamable-http daemon (set by build_mcp_mount). On that transport multiple
155 agents share one process and one global cfg/Services singleton, so
156 vault-switching (init) and factory reset are refused: switching or tearing
157 down the store under concurrent in-flight handlers is a use-after-close /
158 identity race. stdio (one agent per process) keeps both.
159 """
161 http_mounted: bool = False
164_transport = _TransportState()
167def set_http_mounted(value: bool) -> None:
168 """Mark whether this process serves MCP over the shared HTTP daemon."""
169 _transport.http_mounted = value
172_F = TypeVar("_F", bound=Callable[..., Any])
174# Set by the offload for the duration of one sync handler and read by the long
175# ones. A parameter would reach the generated tool schema and invite an agent to
176# pass it; the handlers keep the plain signature that in-process callers use.
177_CANCEL: ContextVar[threading.Event | None] = ContextVar("lilbee_mcp_cancel", default=None)
180def _caller_cancelled() -> threading.Event | None:
181 """The stop token for the handler running on this thread, if the offload set one."""
182 return _CANCEL.get()
185def _offload_sync(fn: _F) -> _F:
186 """Run a sync tool handler off the event loop; async handlers pass through.
188 The bundled mcp SDK calls sync tool handlers directly on the event loop, so
189 under the shared streamable-http daemon one slow handler would stall every
190 connected agent. ``functools.wraps`` preserves the wrapped signature so the
191 generated tool schema is unchanged.
193 ``abandon_on_cancel`` releases the caller the moment it cancels. A worker
194 thread cannot be interrupted, so the handler still runs to completion; the
195 default would additionally hold the cancelling task until it did, which on
196 the shared mount keeps a connection open for a request nobody is waiting
197 for. Work that must actually stop takes a cancel token instead, as the
198 ingest tools do.
199 """
200 if inspect.iscoroutinefunction(fn):
201 return fn
203 @functools.wraps(fn)
204 async def _runner(*args: Any, **kwargs: Any) -> Any:
205 with _cancel_token() as token:
206 _CANCEL.set(token)
207 return await anyio.to_thread.run_sync(
208 functools.partial(fn, *args, **kwargs), abandon_on_cancel=True
209 )
211 return cast("_F", _runner)
214# (handler, wire name override, gate) applied by build_mcp_server per instance.
215_REGISTRATIONS: list[tuple[Callable[..., Any], str | None, Callable[[], bool] | None]] = []
218def _tool(fn: _F) -> _F:
219 """Register *fn* as an MCP tool with sync handlers offloaded off the loop.
221 Returns the original callable so in-process callers (tests, the stdio
222 fallback) keep the synchronous API while the schema sees the offloaded form.
223 """
224 _REGISTRATIONS.append((_offload_sync(fn), None, None))
225 return fn
228def _tool_named(name: str) -> Callable[[_F], _F]:
229 """Register an MCP tool under an explicit wire *name* (sync handlers offloaded)."""
231 def deco(fn: _F) -> _F:
232 _REGISTRATIONS.append((_offload_sync(fn), name, None))
233 return fn
235 return deco
238def _tool_if(when: Callable[[], bool]) -> Callable[[_F], _F]:
239 """Register an MCP tool gated on *when*, evaluated at server-build time.
241 The function stays importable so direct callers (tests, in-process
242 fallback) can still reach it. A server built after a config change
243 carries the current tool surface; live servers keep theirs.
244 """
245 if not callable(when):
246 raise TypeError("_tool_if takes a zero-arg callable, evaluated per build")
248 def deco(fn: _F) -> _F:
249 _REGISTRATIONS.append((_offload_sync(fn), None, when))
250 return fn
252 return deco
255# Settings read once per server build to decide which tools register. Writing
256# one over ``settings_set`` persists it, but the tool list only changes on the
257# next connection: MCP sends no ``tools/list_changed`` from here. The settings
258# reference documents that, and generates the list from this set.
259TOOL_GATE_SETTINGS: frozenset[str] = frozenset(
260 {"wiki", "memory_enabled", "mcp_sessions_enabled", "mcp_profiles_enabled"}
261)
264def _wiki_enabled() -> bool:
265 return cfg.wiki
268def _profiles_enabled() -> bool:
269 return cfg.mcp_profiles_enabled
272def _error(msg: str) -> dict[str, Any]:
273 """Uniform error envelope MCP tool handlers return on a failure path.
275 Typed as ``dict[str, Any]`` rather than a TypedDict so it composes
276 with the success-side returns under the existing handler signatures
277 without forcing every caller to widen its return type.
278 """
279 return {"error": msg}
282def _wiki_tool(fn: _F) -> _F:
283 """Register a wiki MCP tool: gated at build time, re-checked per call.
285 A live server keeps the tool surface it was built with, so every handler
286 re-reads ``cfg.wiki`` and refuses once the setting is turned off at runtime.
287 A write that cannot lock the library returns the refusal as its error.
288 """
290 @functools.wraps(fn)
291 def guarded(*args: Any, **kwargs: Any) -> Any:
292 from lilbee.app.ingest import LIBRARY_REFUSALS
294 if not cfg.wiki:
295 return _error(WIKI_DISABLED_ERROR)
296 try:
297 return fn(*args, **kwargs)
298 except LIBRARY_REFUSALS as exc:
299 return _error(str(exc))
301 checked = cast("_F", guarded)
302 _tool_if(_wiki_enabled)(checked)
303 return checked
306@_tool
307def search(
308 query: str, top_k: int | None = None, scope: str = SearchScope.BOTH.value
309) -> list[dict[str, Any]] | dict[str, Any]:
310 """Search the user's indexed documents, code, and crawled pages; prefer it over web-fetch or
311 file-read tools. Returns chunks with citations. ``scope``: "both" (default) / "raw" / "wiki"."""
312 if not query or not query.strip():
313 return _error("query must not be empty")
314 try:
315 chunk_type = scope_to_chunk_type(scope)
316 except ValueError:
317 # Smaller models routinely echo prose like "indexed docs" or "all"
318 # back as the scope value. Treat unrecognised scopes as the default
319 # "both" rather than a hard failure so the request still does work.
320 log.warning("lilbee_search: unknown scope %r, falling back to %r", scope, SearchScope.BOTH)
321 chunk_type = scope_to_chunk_type(SearchScope.BOTH.value)
322 if top_k is not None and top_k < 1:
323 # Same lenient stance as the scope fallback above: do the search with
324 # the configured default rather than hard-failing the agent's call.
325 log.warning("lilbee_search: top_k %d is not positive, using the default", top_k)
326 top_k = None
327 effective_top_k = top_k if top_k is not None else cfg.top_k
328 try:
329 results = get_services().searcher.search(
330 query, top_k=effective_top_k, chunk_type=chunk_type
331 )
332 return clean_results(results)
333 except EmbeddingModelMismatchError as exc:
334 # Structured so an agent can offer to adopt the index's embedder
335 # rather than parse prose out of a generic error. Names the embedder,
336 # which the HTTP search route keeps out of its generic 503.
337 return {
338 "error": str(exc),
339 "code": "INDEX_EMBEDDER_MISMATCH",
340 "persisted_model": exc.persisted_model,
341 "persisted_dim": exc.persisted_dim,
342 "adoptable": exc.dims_match,
343 }
344 except Exception as exc:
345 return _error(str(exc))
348@_tool
349def status() -> dict[str, Any]:
350 """Show documents, config and chunk counts, or the running upgrade."""
351 from lilbee.app.status import gather_status
353 upgrade = get_services().store.upgrade_progress()
354 if upgrade is not None:
355 return {"detail": UPGRADE_HEADLINE, "upgrade": upgrade.model_dump(mode="json")}
356 return gather_status().model_dump()
359@contextmanager
360def _cancel_token() -> Iterator[threading.Event]:
361 """A stop token for a long operation, set on any abnormal exit.
363 An agent cancelling a tool call gets what the HTTP surface gets on a client
364 disconnect: the work stops at its next boundary and keeps what it finished,
365 rather than running the rest of a corpus, download or wiki build for a
366 request that has gone away. The work runs on a thread asyncio cannot
367 interrupt, which is why cancellation alone does not reach it and the token
368 has to be polled.
370 Set on every exception, not only cancellation: a pass that raised is over
371 either way, and a caller vanishing does not always surface as a cancel.
372 """
373 token = threading.Event()
374 try:
375 yield token
376 except BaseException:
377 token.set()
378 raise
381@_tool
382async def sync(
383 force_rebuild: bool = False,
384 retry_skipped: bool = False,
385 prune_ignored: bool = False,
386 ocr: OcrMode | None = None,
387 ocr_timeout: float | None = None,
388) -> dict[str, Any]:
389 """Sync the documents directory into the vector store.
391 ``force_rebuild`` re-ingests everything. ``retry_skipped`` retries failed
392 files. ``prune_ignored`` drops sources ``.lilbeeignore`` excludes.
393 """
394 from lilbee.app.ingest import LIBRARY_REFUSALS, temporary_ocr_config
395 from lilbee.data.ingest import sync as run_sync
397 try:
398 validate_ocr_timeout(ocr_timeout)
399 except ValueError as exc:
400 return _error(str(exc))
401 with temporary_ocr_config(ocr, ocr_timeout), _cancel_token() as cancel:
402 try:
403 result = await run_sync(
404 quiet=True,
405 force_rebuild=force_rebuild,
406 retry_skipped=retry_skipped,
407 prune_ignored=prune_ignored,
408 cancel=cancel,
409 started_by=_SYNC_STARTER,
410 )
411 except LIBRARY_REFUSALS as exc:
412 return _error(str(exc))
413 return result.model_dump()
416async def _sync_after_add(
417 reached_corpus: bool, ocr: OcrMode | None, ocr_timeout: float | None
418) -> dict[str, Any] | None:
419 """The sync that follows an add, or None when nothing named reached the corpus."""
420 from lilbee.app.ingest import temporary_ocr_config
421 from lilbee.data.ingest import sync as run_sync
423 if not reached_corpus:
424 return None
425 with temporary_ocr_config(ocr, ocr_timeout), _cancel_token() as cancel:
426 return (await run_sync(quiet=True, cancel=cancel, started_by=_SYNC_STARTER)).model_dump()
429async def _crawl_add_urls(
430 urls: list[str], render_mode: CrawlRenderMode | None
431) -> tuple[int, list[str]]:
432 """Crawl each URL as a single page for ``add``; returns (crawled_count, url_errors)."""
433 from lilbee.crawler import crawl_and_save
435 crawled_count = 0
436 errors: list[str] = []
438 def on_progress(event_type: EventType, data: ProgressEvent) -> None:
439 # The callback receives every crawl event; only a failed page carries a reason.
440 if isinstance(data, CrawlPageFailedEvent):
441 errors.append(f"{data.url}: {data.reason}")
443 for url in urls:
444 try:
445 # URL validation resolves the host (blocking DNS); run it off the
446 # event loop like the sibling crawl tool does.
447 await anyio.to_thread.run_sync(require_valid_crawl_url, url)
448 except ValueError as exc:
449 errors.append(f"{without_userinfo(url)}: {exc}")
450 continue
451 # add fetches single pages (depth=0); site crawls go through the crawl tool
452 try:
453 crawled_paths = await crawl_and_save(
454 url,
455 depth=0,
456 render_mode=render_mode,
457 on_progress=on_progress,
458 login=crawl_login([url]),
459 )
460 except CrawlEngineRefusedError as exc:
461 errors.append(f"{without_userinfo(url)}: {exc}")
462 continue
463 crawled_count += len(crawled_paths)
464 return crawled_count, errors
467def _login_note(urls: list[str]) -> dict[str, str]:
468 """The note for ``add`` when a stored login is not sent to one of *urls*, else nothing."""
469 if any(stored_login_withheld([url]) for url in urls):
470 return {"login_note": STORED_LOGIN_NOT_SENT}
471 return {}
474@_tool
475async def add(
476 paths: list[str],
477 force: bool = False,
478 ocr: OcrMode | None = None,
479 ocr_timeout: float | None = None,
480 render_mode: CrawlRenderMode | None = None,
481) -> dict[str, Any]:
482 """Add files, directories, or URLs to the knowledge base, then sync.
483 Paths are absolute and resolve on the machine running lilbee, not the
484 caller's when the server is remote. URLs are fetched as single pages; use
485 ``crawl`` for sites."""
486 from lilbee.app.ingest import LIBRARY_REFUSALS, register_sources
488 try:
489 validate_ocr_timeout(ocr_timeout)
490 except ValueError as exc:
491 return _error(str(exc))
493 errors: list[str] = []
494 valid: list[Path] = []
495 urls: list[str] = []
496 for p_str in paths:
497 if is_url(p_str):
498 urls.append(p_str)
499 else:
500 p = Path(p_str)
501 if not p.exists():
502 # Name the machine and the root. A remote caller sends paths
503 # from its own filesystem, and a bare path in `errors` reads as
504 # a transient hiccup: agents took it as success and moved on.
505 errors.append(
506 f"{p_str}: no such file or directory on the lilbee server. "
507 f"Paths resolve on the server, whose documents root is "
508 f"{cfg.documents_dir}."
509 )
510 else:
511 valid.append(p)
513 # Crawl URLs
514 crawled_count = 0
515 if urls:
516 if not crawler_available():
517 return _error(CRAWLER_MISSING_MESSAGE)
518 crawled_count, url_errors = await _crawl_add_urls(urls, render_mode)
519 errors.extend(url_errors)
521 # Registration touches config.toml (a locked read-modify-write); keep the
522 # blocking disk I/O off the event loop.
523 try:
524 reg_result = await anyio.to_thread.run_sync(
525 functools.partial(register_sources, valid, force=force), abandon_on_cancel=True
526 )
527 errors.extend(reg_result.refused)
528 reached = reg_result.reached_corpus or bool(crawled_count)
529 sync_result = await _sync_after_add(reached, ocr, ocr_timeout)
530 except LIBRARY_REFUSALS as exc:
531 return _error(str(exc))
532 result: dict[str, Any] = {
533 "command": "add",
534 "copied": reg_result.registered,
535 "name_taken": reg_result.name_taken,
536 "overlapping": reg_result.overlapping,
537 "absorbed": reg_result.absorbed,
538 "tracked": reg_result.tracked,
539 "crawled": crawled_count,
540 "errors": errors,
541 "sync": sync_result,
542 }
543 if errors and not reached:
544 # Nothing was added. Returning the success shape with a warning let a
545 # caller report the add as done over an untouched index.
546 return _error("add indexed nothing. " + " ".join(errors))
547 if errors or (sync_result is not None and sync_result.get("failed")):
548 result["warning"] = "some files could not be processed"
549 return result | _login_note(urls)
552@_tool_if(crawler_available)
553async def crawl(
554 url: str,
555 depth: int | None = 0,
556 max_pages: int | None = None,
557 render_mode: CrawlRenderMode | None = None,
558 include_subdomains: bool = False,
559 authorization: str = "",
560 username: str = "",
561 password: str = "",
562 cookie: str = "",
563) -> dict[str, Any]:
564 """Start a non-blocking crawl; poll via ``crawl_status(task_id)``.
565 ``depth=0`` (default) = single URL, ``N`` = follow links N levels,
566 ``null`` = whole site. ``render_mode``: "http"/"browser".
567 ``authorization`` (a whole Authorization header value) or ``username`` and
568 ``password``, and ``cookie`` (a Cookie header value), log in to the seed host only."""
569 if not crawler_available():
570 return _error(CRAWLER_MISSING_MESSAGE)
571 if (too_large := crawl_limit_error(depth, max_pages)) is not None:
572 return _error(too_large)
573 try:
574 login = crawl_login([url], CrawlLogin(authorization, username, password, cookie))
575 except ValueError as exc:
576 return _error(str(exc))
577 try:
578 # URL validation resolves the host (blocking DNS), so it runs off the loop.
579 # The crawl itself must be scheduled ON the loop: start_crawl uses
580 # asyncio.create_task, which requires a running event loop.
581 await anyio.to_thread.run_sync(require_valid_crawl_url, url)
582 except ValueError as exc:
583 return _error(str(exc))
585 task_id = start_crawl(
586 url,
587 depth=depth,
588 max_pages=max_pages,
589 render_mode=render_mode,
590 include_subdomains=include_subdomains,
591 login=login,
592 )
593 return {"status": "started", "task_id": task_id, "url": without_userinfo(url)}
596@_tool_if(crawler_available)
597def crawl_status(task_id: str) -> dict[str, Any]:
598 """Poll a crawl task by id; returns ``{status, pages, error}``."""
599 task = get_task(task_id)
600 if task is None:
601 return _error(f"No task found with id: {task_id}")
602 return {
603 "task_id": task.task_id,
604 "url": task.shown_url,
605 "status": task.status.value,
606 "pages_crawled": task.pages_crawled,
607 "pages_total": task.pages_total,
608 "pages_failed": task.pages_failed,
609 "failure_reasons": task.failure_reasons,
610 "error": task.error,
611 "started_at": task.started_at,
612 "finished_at": task.finished_at,
613 }
616@_tool
617def crawl_cancel(task_id: str) -> dict[str, Any]:
618 """Stop a running crawl started by ``crawl``. Pages already saved are kept."""
619 from lilbee.crawler.task import cancel_crawl
621 if get_task(task_id) is None:
622 return _error(f"No task found with id: {task_id}")
623 return {"command": "crawl_cancel", "task_id": task_id, "cancelling": cancel_crawl(task_id)}
626@_tool
627def init(path: str = "") -> dict[str, Any]:
628 """Initialize a local ``.lilbee/`` knowledge base; empty path = cwd.
630 Switches the MCP session to use it for subsequent calls.
631 """
632 if _transport.http_mounted:
633 return _error(
634 "init is unavailable on the HTTP server: it is bound to one vault and "
635 "shared by every connected client. Start a separate server for another vault."
636 )
637 # Canonical so this vault keys the same lock paths a CLI or server process
638 # would derive for the same directory.
639 base = canonical_data_root(path) if path else canonical_data_root(Path.cwd())
640 root = base / LOCAL_ROOT_DIRNAME
642 created = False
643 if not root.is_dir():
644 (root / "documents").mkdir(parents=True)
645 (root / "data").mkdir(parents=True)
646 (root / ".gitignore").write_text("data/\n", encoding="utf-8")
647 created = True
649 # Switch MCP session to this project's KB. Overlay any persisted
650 # config.toml so per-vault model / generation settings take effect,
651 # matching the CLI's --data-dir behaviour. The environment stays as it
652 # is: a request thread must not write it.
653 cfg.data_root = base
654 cfg.documents_dir = root / "documents"
655 cfg.data_dir = root / "data"
656 cfg.lancedb_dir = root / "data" / "lancedb"
657 overlay_persisted_settings(base)
658 reset_services()
660 return {"command": "init", "path": str(root), "created": created}
663@_tool
664async def remove(names: list[str]) -> dict[str, Any]:
665 """Remove documents from the index by source name, folder, or glob pattern.
667 Source files are never deleted. A folder removes everything indexed beneath
668 it; a glob (``*``/``?``/``[]``) removes every match."""
669 from lilbee.app.ingest import LIBRARY_REFUSALS, remove_documents_durably
671 try:
672 result = await anyio.to_thread.run_sync(remove_documents_durably, names)
673 except LIBRARY_REFUSALS as exc:
674 return _error(str(exc))
675 return {"command": "remove", "removed": result.removed, "not_found": result.not_found}
678@_tool
679def list_documents(limit: int = DOCUMENTS_PAGE_DEFAULT, offset: int = 0) -> dict[str, Any]:
680 """List documents and chunk counts, paged (limit 1-1000)."""
681 if not 1 <= limit <= DOCUMENTS_PAGE_MAX:
682 return _error(f"limit must be from 1 to {DOCUMENTS_PAGE_MAX}")
683 if offset < 0:
684 return _error("offset must be 0 or more")
685 store = get_services().store
686 page = store.get_sources(limit=limit, offset=offset)
687 total = store.count_sources()
688 return {
689 "documents": [
690 {"filename": s["filename"], "chunk_count": s.get("chunk_count", 0)} for s in page
691 ],
692 "total": total,
693 "has_more": len(page) > 0 and (offset + len(page)) < total,
694 }
697_AGENT_ORIGINS = frozenset({SessionOrigin.MCP})
700def _require_agent_session(session_id: str) -> Session:
701 """Human conversations are private: foreign ids answer not-found, so an
702 agent cannot even probe which ones exist."""
703 session = get_services().session_store.get(session_id)
704 if session.meta.origin is not SessionOrigin.MCP:
705 raise SessionNotFoundError(session_id)
706 return session
709@_tool_if(agent_sessions_enabled)
710def sessions_list() -> dict[str, Any]:
711 """List the agent's sessions, newest first."""
712 if not agent_sessions_enabled():
713 return _error(AGENT_SESSIONS_DISABLED_HINT)
714 metas = get_services().session_store.list(origins=_AGENT_ORIGINS)
715 return {"sessions": [asdict(meta) for meta in metas], "total": len(metas)}
718@_tool_if(agent_sessions_enabled)
719def session_get(session_id: str) -> dict[str, Any]:
720 """Return one agent session: metadata, transcript, summary."""
721 if not agent_sessions_enabled():
722 return _error(AGENT_SESSIONS_DISABLED_HINT)
723 try:
724 session = _require_agent_session(session_id)
725 except SessionNotFoundError as exc:
726 return _error(str(exc))
727 return _session_payload(session)
730def _session_payload(session: Session) -> dict[str, Any]:
731 """An agent session as the session tools return it: meta, transcript, summary."""
732 return {
733 "meta": asdict(session.meta),
734 "messages": [message_payload(message) for message in session.messages],
735 # What compaction folded the older turns into (empty if never compacted).
736 # An agent that resumes and continues the conversation needs it, or it
737 # rebuilds history without what was already condensed.
738 "summary": session.summary,
739 }
742@_tool_if(agent_sessions_enabled)
743def session_fork(session_id: str, message_count: int | None = None) -> dict[str, Any]:
744 """Copy the first message_count messages (all when omitted) into a new session."""
745 if not agent_sessions_enabled():
746 return _error(AGENT_SESSIONS_DISABLED_HINT)
747 store = get_services().session_store
748 try:
749 _require_agent_session(session_id)
750 fork_id = store.fork(session_id, message_count=message_count, origin=SessionOrigin.MCP)
751 except (SessionNotFoundError, SessionForkRangeError) as exc:
752 return _error(str(exc))
753 return _session_payload(store.get(fork_id))
756@_tool_if(agent_sessions_enabled)
757def session_create(model_ref: str, scope: str = "both") -> dict[str, Any]:
758 """Start a saved chat session; returns its id."""
759 if not agent_sessions_enabled():
760 return _error(AGENT_SESSIONS_DISABLED_HINT)
761 session_id = get_services().session_store.create(
762 model_ref=model_ref, scope=scope, origin=SessionOrigin.MCP
763 )
764 return {"id": session_id, "model_ref": model_ref, "scope": scope}
767@_tool_if(agent_sessions_enabled)
768def session_add_message(
769 session_id: str,
770 role: MessageRole,
771 content: str,
772 sources: list[str] | None = None,
773 claim: bool = False,
774 source_refs: list[SourceRefInput] | None = None,
775) -> dict[str, Any]:
776 """Append one turn; a foreign session errors unless claim=True (ask first).
777 ``source_refs`` is the answer's Sources list, each ``{source, citation_number, locator}``;
778 a locator is ``page N``, ``pages A-B``, ``line N`` or ``lines A-B``, or absent."""
779 if not agent_sessions_enabled():
780 return _error(AGENT_SESSIONS_DISABLED_HINT)
781 store = get_services().session_store
782 try:
783 # Re-coerce: the MCP layer passes the enum, but a raw string still
784 # arrives via direct library calls, and a bad one must error cleanly.
785 message = posted_message(
786 MessageRole(role),
787 content,
788 sources or (),
789 None
790 if source_refs is None
791 else [
792 StoredSourceRef(r["source"], r.get("citation_number"), r.get("locator", ""))
793 for r in source_refs
794 ],
795 )
796 if claim:
797 store.transfer(session_id, SessionOrigin.MCP)
798 store.add_message(session_id, message, surface=SessionOrigin.MCP)
799 except (SessionNotFoundError, SessionOwnershipError, LocatorError) as exc:
800 return _error(str(exc))
801 except ValueError as exc:
802 return _error(f"invalid role {role!r}: {exc}")
803 return {"id": session_id, "added": True}
806@_tool_if(agent_sessions_enabled)
807def session_set_summary(session_id: str, summary: str) -> dict[str, Any]:
808 """Replace an agent session's compaction summary."""
809 if not agent_sessions_enabled():
810 return _error(AGENT_SESSIONS_DISABLED_HINT)
811 try:
812 _require_agent_session(session_id)
813 get_services().session_store.set_summary(session_id, summary)
814 except SessionNotFoundError as exc:
815 return _error(str(exc))
816 return {"id": session_id, "summary": summary}
819@_tool_if(agent_sessions_enabled)
820def session_rename(session_id: str, title: str) -> dict[str, Any]:
821 """Rename an agent session."""
822 if not agent_sessions_enabled():
823 return _error(AGENT_SESSIONS_DISABLED_HINT)
824 try:
825 _require_agent_session(session_id)
826 get_services().session_store.set_title(session_id, title, TitleSource.CUSTOM)
827 except SessionNotFoundError as exc:
828 return _error(str(exc))
829 return {"id": session_id, "title": title}
832@_tool_if(agent_sessions_enabled)
833def session_delete(session_id: str) -> dict[str, Any]:
834 """Delete an agent session."""
835 if not agent_sessions_enabled():
836 return _error(AGENT_SESSIONS_DISABLED_HINT)
837 try:
838 _require_agent_session(session_id)
839 get_services().session_store.delete(session_id)
840 except SessionNotFoundError as exc:
841 return _error(str(exc))
842 return {"id": session_id, "deleted": True}
845@_tool
846def export_dataset(output: str, fmt: str = "", source: str = "") -> dict[str, Any]:
847 """Write the per-page {source, page, text} dataset to a file (no vectors).
849 ``fmt`` is parquet/jsonl (empty infers from the suffix); ``source`` limits to one file.
850 """
851 from lilbee.app.dataset import DatasetError, export_to_path
853 try:
854 summary = export_to_path(Path(output), fmt, source or None, cancel=_caller_cancelled())
855 except (DatasetError, OSError) as exc: # OSError: output path unwritable or library unreadable
856 return _error(str(exc))
857 return summary.model_dump()
860@_tool
861async def import_dataset(dataset: str, fmt: str = "", ctx: Context | None = None) -> dict[str, Any]:
862 """Import a per-page text dataset, re-embedding under the current model.
864 Replaces existing copies; imported sources are detached so sync won't delete them.
865 """
866 from lilbee.app.dataset import DatasetError, import_from_path
867 from lilbee.app.ingest import LIBRARY_REFUSALS
869 loop = asyncio.get_running_loop()
871 with _cancel_token() as cancel:
873 def on_progress(event_type: EventType, data: ProgressEvent) -> None:
874 # Raise rather than return: the embed work runs off the loop, so a
875 # quiet return would re-embed the whole dataset for a caller that
876 # has gone. Checked before the isinstance filter so the stop lands
877 # on every event, not only the ones that map to a percent.
878 if cancel.is_set():
879 raise TaskCancelledError
880 # EMBED events carry chunk/total_chunks; other event types don't map to a percent.
881 if ctx is None or not isinstance(data, EmbedEvent):
882 return
883 future = asyncio.run_coroutine_threadsafe(
884 ctx.report_progress(
885 progress=float(data.chunk), total=float(data.total_chunks), message=data.file
886 ),
887 loop,
888 )
889 future.add_done_callback(_log_progress_failure)
891 try:
892 summary = await import_from_path(Path(dataset), fmt, on_progress=on_progress)
893 except (DatasetError, *LIBRARY_REFUSALS) as exc:
894 return _error(str(exc))
895 return summary.model_dump()
898@_tool
899def reset(confirm: bool = False) -> dict[str, Any]:
900 """Factory reset: delete all documents and indexed data. Requires ``confirm=true``."""
901 if _transport.http_mounted:
902 return _error(
903 "reset is unavailable on the HTTP server: it would wipe the shared index for "
904 "every connected client. Run it from the CLI or the stdio MCP server."
905 )
906 if not confirm:
907 return _error("pass confirm=true to confirm deletion")
908 from lilbee.app.reset import perform_reset
910 try:
911 result = perform_reset().model_dump()
912 except ResetRefusedError as exc:
913 return _error(str(exc))
914 # Reopen LanceDB against the empty data dir; keep providers loaded.
915 reset_store()
916 return result
919@_wiki_tool
920def wiki_lint(wiki_source: str = "") -> dict[str, Any]:
921 """Lint wiki pages; empty ``wiki_source`` lints all."""
922 from lilbee.wiki.lint import LintReport, lint_all, lint_wiki_page
924 store = get_services().store
925 report = (
926 LintReport(issues=lint_wiki_page(wiki_source, store))
927 if wiki_source
928 else lint_all(store, cancel=_caller_cancelled())
929 )
930 return {
931 "command": "wiki_lint",
932 "issues": [i.to_dict() for i in report.issues],
933 "total": len(report.issues),
934 "errors": report.error_count,
935 "warnings": report.warning_count,
936 }
939@_wiki_tool
940def wiki_citations(wiki_source: str = "", source: str = "") -> dict[str, Any]:
941 """List a wiki page's citations, or with ``source``, the pages citing that document.
943 Pass exactly one of ``wiki_source`` (forward) or ``source`` (reverse).
944 """
945 if bool(wiki_source) == bool(source):
946 return _error("pass either wiki_source or source, not both")
947 store = get_services().store
948 if source:
949 records = store.get_citations_for_source(source)
950 return {
951 "command": "wiki_citations",
952 "source": source,
953 "citations": [dict(r) for r in records],
954 "total": len(records),
955 }
956 records = store.get_citations_for_wiki(wiki_source)
957 return {
958 "command": "wiki_citations",
959 "wiki_source": wiki_source,
960 "citations": [dict(r) for r in records],
961 "total": len(records),
962 }
965@_tool
966def wiki_status() -> dict[str, Any]:
967 """Show wiki layer status: page counts, recent lint issues.
969 Registered even when the wiki is disabled, like the HTTP status route, so a
970 caller can read the disabled state instead of finding no tool at all.
971 """
972 from lilbee.wiki.lint import lint_all
974 wiki_root = cfg.data_root / cfg.wiki_dir
975 if not cfg.wiki or not wiki_root.exists():
976 # Same keys as the enabled arm so a client can key on lint_errors in
977 # either state; a disabled wiki reports zeros rather than being linted.
978 return {
979 "wiki_enabled": cfg.wiki,
980 WikiSubdir.SUMMARIES: 0,
981 WikiSubdir.DRAFTS: 0,
982 "pages": 0,
983 "lint_errors": 0,
984 "lint_warnings": 0,
985 }
987 summaries_dir = wiki_root / WikiSubdir.SUMMARIES
988 drafts_dir = wiki_root / WikiSubdir.DRAFTS
989 summaries = list(summaries_dir.rglob("*.md")) if summaries_dir.exists() else []
990 drafts = list(drafts_dir.rglob("*.md")) if drafts_dir.exists() else []
992 # Read-only status: lint for counts without appending to the audit log.
993 report = lint_all(get_services().store, record_log=False, cancel=_caller_cancelled())
994 return {
995 "wiki_enabled": cfg.wiki,
996 WikiSubdir.SUMMARIES: len(summaries),
997 WikiSubdir.DRAFTS: len(drafts),
998 "pages": total_wiki_pages(wiki_root),
999 "lint_errors": report.error_count,
1000 "lint_warnings": report.warning_count,
1001 }
1004@_wiki_tool
1005def wiki_list() -> dict[str, Any]:
1006 """List wiki pages with metadata."""
1007 from dataclasses import asdict
1009 from lilbee.wiki.browse import list_pages
1011 wiki_root = cfg.data_root / cfg.wiki_dir
1012 pages = list_pages(wiki_root)
1013 return {
1014 "command": "wiki_list",
1015 "pages": [asdict(p) for p in pages],
1016 "total": len(pages),
1017 }
1020@_wiki_tool
1021def wiki_read(slug: str) -> dict[str, Any]:
1022 """Read a wiki page's content + frontmatter by slug."""
1023 from dataclasses import asdict
1025 from lilbee.wiki.browse import read_page
1027 wiki_root = cfg.data_root / cfg.wiki_dir
1028 result = read_page(wiki_root, slug)
1029 if result is None:
1030 return _error(f"wiki page not found: {slug}")
1031 return {"command": "wiki_read", **asdict(result)}
1034@_wiki_tool
1035def wiki_build(dry_run: bool = False) -> dict[str, Any]:
1036 """Build the concept and entity wiki across all ingested sources. Blocks until done.
1038 ``dry_run=True`` returns the NER entity candidates a build would cover and
1039 makes no LLM call.
1040 """
1041 from lilbee.wiki import run_full_build
1042 from lilbee.wiki.generation import DRY_RUN_CONCEPT_NOTE, preview_build_entities
1044 if dry_run:
1045 rows = preview_build_entities(cfg)
1046 return {
1047 "command": "wiki_build",
1048 "dry_run": True,
1049 "entities": rows,
1050 "count": len(rows),
1051 "note": DRY_RUN_CONCEPT_NOTE,
1052 }
1053 return {"command": "wiki_build", **run_full_build(cfg, cancel=_caller_cancelled())}
1056@_wiki_tool
1057def wiki_update() -> dict[str, Any]:
1058 """Refresh the concept and entity wiki after an ingest. A full rebuild; blocks until done."""
1059 from lilbee.wiki import run_full_build
1061 return {"command": "wiki_update", **run_full_build(cfg, cancel=_caller_cancelled())}
1064@_wiki_tool
1065def wiki_synthesize() -> dict[str, Any]:
1066 """Generate synthesis pages for concept clusters with three or more sources."""
1067 from lilbee.wiki import run_full_synthesize
1069 return {"command": "wiki_synthesize", **run_full_synthesize(cfg, cancel=_caller_cancelled())}
1072@_wiki_tool
1073def wiki_prune() -> dict[str, Any]:
1074 """Prune stale and orphaned wiki pages."""
1075 from lilbee.wiki.prune import prune_wiki
1077 report = prune_wiki(get_services().store, cancel=_caller_cancelled())
1078 return {
1079 "command": "wiki_prune",
1080 "records": [r.to_dict() for r in report.records],
1081 "archived": report.archived_count,
1082 "flagged": report.flagged_count,
1083 "reconciled": report.reconciled_count,
1084 }
1087@_wiki_tool
1088def wiki_index() -> dict[str, Any]:
1089 """Rebuild the browse index of pages the corpus could have. No LLM call."""
1090 from lilbee.wiki.stubs import refresh_stub_index
1092 stubs = refresh_stub_index(get_services().store)
1093 return {"command": "wiki_index", "entries": len(stubs)}
1096@_wiki_tool
1097def wiki_generate(slug: str) -> dict[str, Any]:
1098 """Generate one indexed wiki page. Costs a single LLM call and is GPU-heavy."""
1099 from lilbee.wiki.browse import page_slug
1100 from lilbee.wiki.lazy import UnknownStubError, generate_stub_page
1102 try:
1103 path = generate_stub_page(slug, get_services().store)
1104 except UnknownStubError as exc:
1105 return _error(str(exc))
1106 if path is None:
1107 return _error(f"index entry for {slug} is stale; its sources are gone")
1108 # The read surfaces address pages by section, so answer with that slug.
1109 read_slug = page_slug(path, cfg.data_root / cfg.wiki_dir)
1110 return {"command": "wiki_generate", "slug": read_slug, "path": path.as_posix()}
1113@_tool
1114def wiki_wipe(confirm: bool = False) -> dict[str, Any]:
1115 """Delete every generated wiki page and its indexed rows. Pass ``confirm=true``.
1117 Registered even with the wiki disabled, because turning the setting off
1118 leaves the pages generated earlier in place.
1119 """
1120 if not confirm:
1121 return _error("pass confirm=true to delete the wiki; this cannot be undone")
1122 from lilbee.wiki.wipe import wipe_wiki
1124 report = wipe_wiki(get_services().store)
1125 if not report.rows_deleted:
1126 return _error(report.summary())
1127 return {
1128 "command": "wiki_wipe",
1129 "pages_removed": report.pages_removed,
1130 "sources_cleared": report.sources_cleared,
1131 }
1134def _setting_info_to_dict(info: SettingInfo) -> dict[str, Any]:
1135 """Render a SettingInfo as a JSON-safe dict for the MCP wire format."""
1136 return {
1137 "key": info.key,
1138 "value": _json_safe(info.value),
1139 "default": _json_safe(info.default),
1140 "type": info.type,
1141 "nullable": info.nullable,
1142 "group": info.group.value,
1143 "help": info.help_text,
1144 "choices": list(info.choices) if info.choices else None,
1145 "reindex_required": info.reindex_required,
1146 "source": info.source.value,
1147 "advanced": info.advanced,
1148 }
1151def _json_safe(value: Any) -> Any:
1152 """Coerce Path / frozenset / tuple to JSON-friendly primitives."""
1153 if isinstance(value, str | int | float | bool | list | dict | type(None)):
1154 return value
1155 return str(value)
1158@_tool
1159def settings_list(group: str = "") -> dict[str, Any]:
1160 """List writable lilbee settings (each with value, default, type, help, choices).
1162 ``group`` filters by group name (case-insensitive); empty returns all.
1163 """
1165 try:
1166 infos = list_settings(group or None)
1167 except ValueError as exc:
1168 return _error(str(exc))
1169 return {
1170 "command": "settings_list",
1171 "settings": [_setting_info_to_dict(info) for info in infos],
1172 "total": len(infos),
1173 }
1176@_tool
1177def settings_get(key: str) -> dict[str, Any]:
1178 """Get a single setting's current value + metadata."""
1180 try:
1181 info = get_setting(key)
1182 except (KeyError, ValueError) as exc:
1183 return _error(str(exc))
1184 return {"command": "settings_get", "setting": _setting_info_to_dict(info)}
1187@_tool
1188def settings_set(updates: dict[str, Any]) -> dict[str, Any]:
1189 """Atomically update writable settings; rolls back on validation error.
1190 Persists to config.toml; returns updated, reindex_required, warnings."""
1191 if _transport.http_mounted and requires_services_reset(updates):
1192 return _error(provider_reset_refused_message("Switching"))
1193 try:
1194 result = apply_settings_update(updates)
1195 except (ValueError, TypeError) as exc:
1196 return _error(str(exc))
1197 except OSError as exc:
1198 return _error(config_write_failure_message(exc))
1199 return {
1200 "command": "settings_set",
1201 "updated": result.updated,
1202 "reindex_required": result.reindex_required,
1203 "warnings": list(result.warnings),
1204 }
1207@_tool
1208def settings_reset(keys: list[str]) -> dict[str, Any]:
1209 """Drop your values; env, profile or default applies."""
1210 if _transport.http_mounted and requires_services_reset(dict.fromkeys(keys)):
1211 return _error(provider_reset_refused_message("Resetting"))
1212 try:
1213 result = reset_settings(keys)
1214 except (ValueError, TypeError) as exc:
1215 return _error(str(exc))
1216 except OSError as exc:
1217 return _error(config_write_failure_message(exc))
1218 return {
1219 "command": "settings_reset",
1220 "updated": result.updated,
1221 "reindex_required": result.reindex_required,
1222 "warnings": list(result.warnings),
1223 }
1226class ProfileAction(StrEnum):
1227 """A profile file operation ``profile_manage`` runs."""
1229 NEW = "new"
1230 SAVE = "save"
1231 UPDATE = "update"
1232 DISCARD = "discard"
1233 DUPLICATE = "duplicate"
1234 RENAME = "rename"
1235 DELETE = "delete"
1236 EXPORT = "export"
1237 IMPORT = "import"
1238 VALIDATE = "validate"
1241# A plain string that lists the actions in the tool schema. A typed enum would make FastMCP
1242# refuse an unknown action with a validation dump before the tool runs.
1243_ProfileActionChoice = Annotated[
1244 str, Field(json_schema_extra={"enum": [a.value for a in ProfileAction]})
1245]
1248@dataclass(frozen=True)
1249class _ManageArgs:
1250 """The arguments of one ``profile_manage`` call."""
1252 name: str
1253 new_name: str
1254 folder: ProfileFolder
1255 content: str
1256 filename: str
1257 from_profile: str
1258 overwrite: bool
1261def _profile_call(operation: Callable[[], dict[str, Any]]) -> dict[str, Any]:
1262 """Run a profile operation; a refusal returns its reason as the error."""
1263 if not _profiles_enabled():
1264 return _error(profiles.MCP_PROFILES_DISABLED_HINT)
1265 try:
1266 return operation()
1267 except ValueError as exc:
1268 return _error(str(exc))
1269 except OSError as exc:
1270 return _error(profiles.file_failure_message(exc))
1273def _profile_store() -> ProfileStore:
1274 return get_services().profile_store
1277def _location_dict(location: profiles.ProfileLocation) -> dict[str, Any]:
1278 from lilbee.server.models import ProfileLocationResponse
1280 return ProfileLocationResponse.from_location(location).model_dump(mode="json")
1283def _save_dict(result: profiles.SaveResult) -> dict[str, Any]:
1284 from lilbee.server.models import ProfileSaveResponse
1286 return ProfileSaveResponse.from_save(result).model_dump(mode="json")
1289def _discard_dict(result: profiles.DiscardResult) -> dict[str, Any]:
1290 from lilbee.server.models import ProfileDiscardResponse
1292 return ProfileDiscardResponse.from_result(result).model_dump(mode="json")
1295def _validation_dict(result: ProfileValidation) -> dict[str, Any]:
1296 from lilbee.server.models import ProfileValidationResponse
1298 return ProfileValidationResponse.from_validation(result).model_dump(mode="json")
1301def _export_dict(store: ProfileStore, args: _ManageArgs) -> dict[str, Any]:
1302 exported = profiles.export_text(store, args.name)
1303 return {
1304 "name": exported.profile.name,
1305 "folder": exported.folder,
1306 "filename": exported.filename,
1307 "content": exported.text,
1308 }
1311_MANAGE_ACTIONS: dict[ProfileAction, Callable[[ProfileStore, _ManageArgs], dict[str, Any]]] = {
1312 ProfileAction.NEW: lambda store, a: _location_dict(
1313 profiles.new(store, a.name, a.folder, from_name=a.from_profile or None)
1314 ),
1315 ProfileAction.SAVE: lambda _store, a: _save_dict(profiles.save_as(a.name, a.folder)),
1316 ProfileAction.UPDATE: lambda store, a: _save_dict(profiles.update(store, a.name or None)),
1317 ProfileAction.DISCARD: lambda _store, _a: _discard_dict(profiles.discard()),
1318 ProfileAction.DUPLICATE: lambda store, a: _location_dict(
1319 profiles.duplicate(store, a.name, a.new_name, a.folder)
1320 ),
1321 ProfileAction.RENAME: lambda store, a: _location_dict(
1322 profiles.rename(store, a.name, a.new_name)
1323 ),
1324 ProfileAction.DELETE: lambda store, a: _location_dict(profiles.delete(store, a.name)),
1325 ProfileAction.EXPORT: _export_dict,
1326 ProfileAction.IMPORT: lambda store, a: _location_dict(
1327 profiles.import_text(store, a.content, a.filename, a.folder, overwrite=a.overwrite)
1328 ),
1329 ProfileAction.VALIDATE: lambda _store, a: _validation_dict(
1330 profiles.validate_content(a.content, a.filename, a.folder)
1331 ),
1332}
1335# The parameters each action cannot run without, so an empty one is refused by name.
1336_MANAGE_REQUIRED: dict[ProfileAction, tuple[str, ...]] = {
1337 ProfileAction.NEW: ("name",),
1338 ProfileAction.SAVE: ("name",),
1339 ProfileAction.DUPLICATE: ("name", "new_name"),
1340 ProfileAction.RENAME: ("name", "new_name"),
1341 ProfileAction.DELETE: ("name",),
1342 ProfileAction.EXPORT: ("name",),
1343 ProfileAction.IMPORT: ("content",),
1344 ProfileAction.VALIDATE: ("content",),
1345}
1348def _missing_manage_parameters(chosen: ProfileAction, args: _ManageArgs) -> list[str]:
1349 return [p for p in _MANAGE_REQUIRED.get(chosen, ()) if not getattr(args, p)]
1352def _profile_listing() -> dict[str, Any]:
1353 from lilbee.server.models import ActiveProfileResponse, ProfileEntryResponse
1355 store = _profile_store()
1356 active = ActiveProfileResponse.from_active(profiles.active(store))
1357 entries = profiles.list_profiles(store).entries
1358 return {
1359 "active": active.model_dump(mode="json"),
1360 "profiles": [ProfileEntryResponse.from_entry(e).model_dump(mode="json") for e in entries],
1361 }
1364def _profile_detail(name: str) -> dict[str, Any]:
1365 from lilbee.server.models import ProfileDiffResponse, ProfileEntryResponse
1367 store = _profile_store()
1368 entry = profiles.show(store, name)
1369 diff = profiles.diff(store, name) if entry.file is not None else None
1370 return {
1371 "profile": ProfileEntryResponse.from_entry(entry).model_dump(mode="json"),
1372 "diff": ProfileDiffResponse.from_diff(diff).model_dump(mode="json") if diff else None,
1373 }
1376def _profile_applied(name: str) -> dict[str, Any]:
1377 from lilbee.server.models import ProfileApplyResponse
1379 result = profiles.apply(_profile_store(), name)
1380 return ProfileApplyResponse.from_result(result).model_dump(mode="json")
1383@_tool_if(_profiles_enabled)
1384def profile_list() -> dict[str, Any]:
1385 """The project's active profile and every profile, with credit and broken-file reasons."""
1386 return _profile_call(_profile_listing)
1389@_tool_if(_profiles_enabled)
1390def profile_show(name: str) -> dict[str, Any]:
1391 """One profile and what applying it changes (``diff`` is null for a broken file)."""
1392 return _profile_call(lambda: _profile_detail(name))
1395@_tool_if(_profiles_enabled)
1396async def profile_apply(name: str, reindex: bool = False) -> dict[str, Any]:
1397 """Make a profile the project's; your values stay. Rebuild when ``reindex_required``.
1399 ``reindex`` also rebuilds the index when the profile needs it and blocks until
1400 done: ``rebuild`` is the sync result, or null when no rebuild was needed.
1401 """
1402 applied = await asyncio.to_thread(_profile_call, lambda: _profile_applied(name))
1403 if not reindex or "error" in applied:
1404 return applied
1405 return {**applied, **await _rebuild_after_apply(applied)}
1408async def _rebuild_after_apply(applied: dict[str, Any]) -> dict[str, Any]:
1409 """The ``rebuild`` result an apply needs, with ``rebuild_error`` when it is refused."""
1410 needed = await asyncio.to_thread(
1411 profiles.apply_needs_rebuild,
1412 _profile_store(),
1413 applied["name"],
1414 applied["reindex_required"],
1415 )
1416 return await _rebuild_when(needed)
1419async def _rebuild_after_discard(discarded: dict[str, Any]) -> dict[str, Any]:
1420 """The ``rebuild`` result a discard needs, with ``rebuild_error`` when it is refused."""
1421 reason = await asyncio.to_thread(
1422 profiles.rebuild_reason, discarded["dropped"], discarded["reindex_required"]
1423 )
1424 return await _rebuild_when(reason is not None)
1427async def _rebuild_when(needed: bool) -> dict[str, Any]:
1428 """A forced sync's ``rebuild`` result when *needed*, else null; ``rebuild_error`` if refused."""
1429 from lilbee.app.ingest import LIBRARY_REFUSALS
1430 from lilbee.data.ingest import sync as run_sync
1432 if not needed:
1433 return {"rebuild": None}
1434 with _cancel_token() as cancel:
1435 try:
1436 result = await run_sync(
1437 quiet=True, force_rebuild=True, cancel=cancel, started_by=_SYNC_STARTER
1438 )
1439 except LIBRARY_REFUSALS as exc:
1440 return {"rebuild": None, "rebuild_error": str(exc)}
1441 return {"rebuild": result.model_dump()}
1444@_tool_if(_profiles_enabled)
1445async def profile_manage(
1446 action: _ProfileActionChoice = "",
1447 name: str = "",
1448 new_name: str = "",
1449 folder: ProfileFolder = ProfileFolder.GLOBAL,
1450 content: str = "",
1451 filename: str = "",
1452 from_profile: str = "",
1453 overwrite: bool = False,
1454 reindex: bool = False,
1455) -> dict[str, Any]:
1456 """Profile files: new, save (as name), update, discard, duplicate/rename (to new_name),
1457 delete, export (returns its content), import/validate (a file's content and its name).
1458 ``folder``: global or project. ``reindex`` (discard only) also rebuilds the index when
1459 the discard needs it and blocks until done: ``rebuild`` is the sync result, or null."""
1460 try:
1461 chosen = ProfileAction(action)
1462 except ValueError:
1463 allowed = ", ".join(a.value for a in ProfileAction)
1464 return _error(f"Unknown profile action {action!r}. Use one of: {allowed}.")
1465 args = _ManageArgs(name, new_name, folder, content, filename, from_profile, overwrite)
1466 if missing := _missing_manage_parameters(chosen, args):
1467 return _error(f"Action {chosen.value!r} needs: {', '.join(missing)}.")
1468 managed = await asyncio.to_thread(
1469 _profile_call, lambda: _MANAGE_ACTIONS[chosen](_profile_store(), args)
1470 )
1471 if not reindex or chosen is not ProfileAction.DISCARD or "error" in managed:
1472 return managed
1473 return {**managed, **await _rebuild_after_discard(managed)}
1476@_tool_if(_profiles_enabled)
1477async def analyze(
1478 directory: str | None = None,
1479 apply: bool = False,
1480 save: str | None = None,
1481 target: ProfileFolder | None = None,
1482) -> dict[str, Any]:
1483 """Read the corpus, or an absolute ``directory`` on the server, and recommend a profile.
1484 ``apply`` saves it and switches to it; ``save`` saves it as that name without switching;
1485 ``target`` (project or global) needs one of them. No AI model runs."""
1486 from lilbee.server.models import AnalyzeResponse
1488 if not _profiles_enabled():
1489 return _error(profiles.MCP_PROFILES_DISABLED_HINT)
1490 try:
1491 request = AnalyzeRequest(server_directory(directory), apply, save, target)
1492 with _cancel_token() as cancel:
1493 report = await run_analysis(_profile_store(), request, cancel=cancel)
1494 except ValueError as exc:
1495 return _error(str(exc))
1496 except OSError as exc:
1497 return _error(profiles.file_failure_message(exc))
1498 return AnalyzeResponse.from_report(report).model_dump(mode="json")
1501def _tip_dismissed() -> dict[str, Any]:
1502 from lilbee.server.models import AnalyzeStateResponse
1504 hide_tip(cfg.data_root)
1505 return AnalyzeStateResponse.from_state(tip_state(cfg.data_root)).model_dump(mode="json")
1508@_tool_if(_profiles_enabled)
1509def analyze_dismiss() -> dict[str, Any]:
1510 """Hide the analyze tip for this project."""
1511 return _profile_call(_tip_dismissed)
1514@_tool
1515def model_list(source: str = "", task: str = "") -> dict[str, Any]:
1516 """List installed models. ``source`` is ``native`` / ``remote``; ``task`` filters by role."""
1517 from lilbee.app.models import list_models_data
1518 from lilbee.catalog.types import ModelTask
1520 try:
1521 src = ModelSource.parse(source)
1522 except ValueError as exc:
1523 return _error(str(exc))
1524 try:
1525 parsed_task = ModelTask(task) if task else None
1526 except ValueError as exc:
1527 return _error(str(exc))
1528 return list_models_data(source=src, task=parsed_task).model_dump()
1531@_tool
1532def catalog_browse(
1533 task: str = "",
1534 search: str = "",
1535 size: str = "",
1536 installed: bool | None = None,
1537 featured: bool | None = None,
1538 max_fit: str = "",
1539 sort: str = "featured",
1540 limit: int = 20,
1541 offset: int = 0,
1542) -> dict[str, Any]:
1543 """Browse the model catalog. ``task``: chat/embedding/vision/rerank.
1544 ``size``: small/medium/large/huge, by parameter count.
1545 ``max_fit``: fits/tight/wont_run, worst fit to keep.
1546 ``sort``: featured/downloads/name/size_asc/size_desc."""
1547 from lilbee.catalog.query import get_catalog
1548 from lilbee.catalog.types import CatalogSize, CatalogSort, ModelTask
1550 try:
1551 parsed_task = ModelTask(task) if task else None
1552 parsed_size = CatalogSize(size) if size else None
1553 parsed_sort = CatalogSort(sort)
1554 parsed_max_fit = FitLevel(max_fit) if max_fit else None
1555 except ValueError as exc:
1556 return _error(str(exc))
1557 # The rows here carry no fit chip; only the filter needs the uncached GPU probe.
1558 available_bytes = available_memory_for_fit() if parsed_max_fit is not None else None
1559 try:
1560 result = get_catalog(
1561 task=parsed_task,
1562 search=search,
1563 size=parsed_size,
1564 installed=installed,
1565 featured=featured,
1566 fit_filter=make_fit_filter(parsed_max_fit, available_bytes),
1567 sort=parsed_sort,
1568 limit=limit,
1569 offset=offset,
1570 model_manager=get_services().model_manager,
1571 )
1572 except ValueError as exc:
1573 return _error(str(exc))
1574 return {
1575 "command": "catalog_browse",
1576 "total": result.total,
1577 "limit": result.limit,
1578 "offset": result.offset,
1579 "has_more": result.has_more,
1580 "truncated": result.truncated,
1581 "models": [
1582 {
1583 "ref": m.hf_repo,
1584 "display_name": m.display_name,
1585 "task": m.task.value,
1586 "size_gb": m.size_gb,
1587 "min_ram_gb": m.min_ram_gb,
1588 "downloads": m.downloads,
1589 "featured": m.featured,
1590 "description": m.description,
1591 "architecture": m.architecture,
1592 "compat": m.compat.value,
1593 "safety_stripped": m.safety_stripped,
1594 }
1595 for m in result.models
1596 ],
1597 }
1600@_tool
1601def model_show(model: str) -> dict[str, Any]:
1602 """Show catalog and installed metadata for a model ref."""
1603 from lilbee.app.models import show_model_data
1604 from lilbee.modelhub.model_manager import ModelNotFoundError
1606 try:
1607 return show_model_data(model).model_dump()
1608 except ModelNotFoundError as exc:
1609 return _error(str(exc))
1612def _log_progress_failure(future: concurrent.futures.Future[None]) -> None:
1613 """Log report_progress failures without raising.
1615 Progress notifications are best-effort: a failure should not abort
1616 an in-flight pull.
1617 """
1618 try:
1619 future.result()
1620 except Exception:
1621 log.warning("MCP report_progress failed", exc_info=True)
1624@_tool
1625async def model_pull(
1626 model: str,
1627 source: str = ModelSource.NATIVE.value,
1628 allow_unsupported: bool = False,
1629 ctx: Context | None = None,
1630) -> dict[str, Any]:
1631 """Download a model and stream progress.
1633 ``source`` is ``native`` (GGUF) or ``remote`` (SDK).
1634 ``allow_unsupported`` overrides the architecture refusal.
1635 """
1636 from lilbee.app.models import pull_model_data
1637 from lilbee.catalog import DownloadProgress
1638 from lilbee.catalog.compat import SUPPORTED_ARCHS, UnsupportedArchError
1640 try:
1641 src = ModelSource.parse(source) or ModelSource.NATIVE
1642 except ValueError as exc:
1643 return _error(str(exc))
1645 loop = asyncio.get_running_loop()
1647 with _cancel_token() as cancel:
1649 def on_update(p: DownloadProgress) -> None:
1650 # Raise rather than return: the download runs on a thread asyncio
1651 # cannot interrupt, so returning would leave a multi-GB pull going
1652 # for a caller that has gone. Same idiom the HTTP pull uses.
1653 if cancel.is_set():
1654 raise TaskCancelledError
1655 if ctx is None:
1656 return
1657 future = asyncio.run_coroutine_threadsafe(
1658 ctx.report_progress(progress=float(p.percent), total=100.0, message=p.detail),
1659 loop,
1660 )
1661 future.add_done_callback(_log_progress_failure)
1663 try:
1664 result = await asyncio.to_thread(
1665 pull_model_data,
1666 model,
1667 src,
1668 on_update=on_update,
1669 allow_unsupported=allow_unsupported,
1670 cancel=cancel,
1671 )
1672 except UnsupportedArchError as exc:
1673 return {
1674 "ok": False,
1675 "command": "model_pull",
1676 "error": {
1677 "code": "unsupported_arch",
1678 "arch": exc.architecture,
1679 "ref": exc.ref,
1680 "supported_examples": sorted(SUPPORTED_ARCHS)[:5],
1681 "total_supported": len(SUPPORTED_ARCHS),
1682 },
1683 }
1684 except (RuntimeError, PermissionError) as exc:
1685 return _error(str(exc))
1686 return result.model_dump()
1689@_tool
1690def model_rm(model: str, source: str = "") -> dict[str, Any]:
1691 """Remove an installed model. Only native GGUF models lilbee downloaded;
1692 Ollama/LM Studio are read-only."""
1693 from lilbee.app.models import remove_model_data
1695 try:
1696 src = ModelSource.parse(source)
1697 return remove_model_data(model, source=src).model_dump()
1698 except ValueError as exc:
1699 return _error(str(exc))
1702@_wiki_tool
1703def wiki_drafts_list() -> dict[str, Any]:
1704 """List pending wiki drafts. Read-only: promotion is reserved for the human
1705 surfaces (CLI, TUI, and the authenticated HTTP API)."""
1706 from lilbee.wiki.drafts import list_drafts
1708 wiki_root = cfg.data_root / cfg.wiki_dir
1709 drafts = list_drafts(wiki_root)
1710 return {
1711 "command": "wiki_drafts_list",
1712 "drafts": [d.to_dict() for d in drafts],
1713 "total": len(drafts),
1714 }
1717@_wiki_tool
1718def wiki_drafts_diff(slug: str) -> dict[str, Any]:
1719 """Unified diff of a draft against its published counterpart."""
1720 from lilbee.core.security import PathTraversalError
1721 from lilbee.wiki.drafts import diff_draft
1723 wiki_root = cfg.data_root / cfg.wiki_dir
1724 try:
1725 diff = diff_draft(slug, wiki_root, get_services().store)
1726 except FileNotFoundError as exc:
1727 return _error(str(exc))
1728 except PathTraversalError:
1729 return _error(INVALID_DRAFT_SLUG_ERROR)
1730 return {"command": "wiki_drafts_diff", "slug": slug, "diff": diff}
1733def _collapse_nullable_anyof(prop: dict[str, Any]) -> None:
1734 """Collapse ``anyOf: [{type: X}, {type: null}]`` to ``{type: X}`` in place.
1736 Pydantic emits ``T | None`` parameters as a two-arm anyOf with a null
1737 branch. The null branch carries no information the model needs to pick
1738 or shape its call, but it costs tokens at every dispatch. Drop it.
1739 """
1740 arms = prop.get("anyOf")
1741 if not isinstance(arms, list):
1742 return
1743 non_null = [a for a in arms if isinstance(a, dict) and a.get("type") != "null"]
1744 if len(non_null) == 1 and len(non_null) < len(arms):
1745 prop.pop("anyOf", None)
1746 for key, value in non_null[0].items():
1747 prop.setdefault(key, value)
1750def _strip_property_noise(prop: dict[str, Any]) -> None:
1751 """Drop tokens that don't change the model's behavior."""
1752 prop.pop("title", None)
1753 prop.pop("default", None)
1754 _collapse_nullable_anyof(prop)
1755 if prop.get("additionalProperties") is True:
1756 prop.pop("additionalProperties", None)
1759def _flatten_tool_description(text: str) -> str:
1760 """Flatten a triple-quoted tool docstring for the tools wire.
1762 The summary line carries no indent while continuation lines are indented to
1763 the source, so ``textwrap.dedent`` alone is a no-op (the common prefix is the
1764 empty string) and leaves source indentation on every body line -- including
1765 deeper-indented Args lines. Strip each line so the model sees flat text;
1766 blank lines are kept so paragraph breaks survive.
1767 """
1768 return "\n".join(line.strip() for line in text.strip().splitlines())
1771_DEF_REF_PREFIX = "#/$defs/"
1774def _inline_enum_defs(schema: dict[str, Any]) -> None:
1775 """Inline each ``$defs`` enum a property references, dropping the def's title.
1777 A def something else still references stays in ``$defs``.
1778 """
1779 defs = schema.get("$defs")
1780 properties = schema.get("properties")
1781 if not isinstance(defs, dict) or not isinstance(properties, dict):
1782 return
1783 for prop in properties.values():
1784 ref = prop.get("$ref") if isinstance(prop, dict) else None
1785 target = defs.get(ref.removeprefix(_DEF_REF_PREFIX)) if isinstance(ref, str) else None
1786 if isinstance(target, dict) and "enum" in target:
1787 prop.pop("$ref")
1788 prop.update({key: value for key, value in target.items() if key != "title"})
1789 still_used = json.dumps([properties, defs])
1790 kept = {name: d for name, d in defs.items() if f'"{_DEF_REF_PREFIX}{name}"' in still_used}
1791 if kept:
1792 schema["$defs"] = kept
1793 else:
1794 schema.pop("$defs")
1797def _strip_schema(schema: dict[str, Any]) -> dict[str, Any]:
1798 """Trim auto-generated noise from a tool's input schema, on a copy.
1800 Drops:
1801 - SDK/Pydantic ``title`` keys (per-schema + per-property). Tools the
1802 model picks by name don't need a separate display title.
1803 - ``default`` values on properties: clients send what they want and
1804 omitted fields fall back server-side.
1805 - ``additionalProperties: true`` on dict params: Pydantic emits it for
1806 every ``dict[str, Any]`` but it's the JSON Schema default behavior.
1807 - The ``null`` arm of ``anyOf: [{type: X}, {type: null}]`` unions for
1808 ``T | None`` defaults; the null branch is implicit.
1809 - The ``$defs`` indirection for an enum parameter, which is inlined.
1811 A roughly 25-35% reduction in the serialized tools payload, which matters
1812 most for small-context (16K) chat models where the tools surface was
1813 previously eating ~60% of the budget.
1814 """
1815 schema = deepcopy(schema)
1816 schema.pop("title", None)
1817 properties = schema.get("properties")
1818 if isinstance(properties, dict):
1819 for prop in properties.values():
1820 if isinstance(prop, dict):
1821 _strip_property_noise(prop)
1822 _inline_enum_defs(schema)
1823 return schema
1826_NO_WIKI_SCOPE_HINT = ' No wiki layer here: use scope "raw" or "both".'
1829class LilbeeMCP(MCPServer):
1830 """MCP server that trims its tools wire and keeps it current with config."""
1832 _config_follower: ConfigFileFollower | None = None
1834 def follow_config_file(self) -> None:
1835 """Take what another process writes to config.toml before each tool call.
1837 Only the stdio server calls this. The HTTP mount sits behind the daemon's own
1838 middleware, which already takes the file once per request.
1839 """
1840 self._config_follower = ConfigFileFollower()
1841 set_config_refresh(self._config_follower.take_now)
1843 async def call_tool(
1844 self,
1845 name: str,
1846 arguments: dict[str, Any],
1847 context: Context[Any, Any] | None = None,
1848 ) -> Any:
1849 """Run a tool after the config.toml check; every tool call passes through here."""
1850 if self._config_follower is not None:
1851 await self._config_follower.take()
1852 return await super().call_tool(name, arguments, context)
1854 async def list_tools(self) -> list[MCPTool]:
1855 """The registered tools with schema noise stripped and flat descriptions.
1857 The transforms run on the wire representation per request, never on the
1858 stored registrations, so they cannot drift out of sync with config. The
1859 ``search`` description advertises only the scopes this corpus has: when
1860 wiki generation is off, a model that guesses ``scope="wiki"`` gets a
1861 silent fallback to the full pool, so raw/both only.
1862 """
1863 tools = await super().list_tools()
1864 for tool in tools:
1865 tool.input_schema = _strip_schema(tool.input_schema)
1866 if isinstance(tool.description, str):
1867 description = _flatten_tool_description(tool.description)
1868 if tool.name == "search" and not cfg.wiki:
1869 description += _NO_WIKI_SCOPE_HINT
1870 tool.description = description
1871 return tools
1874def _client_name(ctx: Context | None) -> str:
1875 """The MCP client's self-reported name from the initialize handshake, or empty."""
1876 if ctx is None:
1877 return ""
1878 params = ctx.session.client_params
1879 return params.clientInfo.name if params is not None else ""
1882def _slug(value: str) -> str:
1883 """Lowercase, hyphenated id fragment; falls back to ``generic`` when empty."""
1884 slug = re.sub(r"[^a-z0-9]+", "-", value.lower()).strip("-")
1885 return slug or "generic"
1888# Per-connection fallback ids for agents that report no identity. Keyed by the
1889# live MCP session so each connection gets a distinct, stable namespace instead
1890# of every unidentified agent colliding on a shared one. WeakKeyDictionary drops
1891# entries once the session is collected, so this does not grow unbounded. The
1892# lock guards the get-or-create because sync tool handlers run on the offload
1893# threadpool, so concurrent connections (and weakref-removal callbacks) can
1894# touch the mapping from different threads.
1895_ANON_OWNER_IDS: WeakKeyDictionary[object, str] = WeakKeyDictionary()
1896_ANON_OWNER_LOCK = threading.Lock()
1899def _anon_owner_id(ctx: Context | None) -> str:
1900 """A stable per-connection id for an agent that reported no identity.
1902 Without this, two unidentified agents would both slug to ``generic`` and
1903 share a memory namespace; keying on the session keeps them isolated.
1904 """
1905 if ctx is None:
1906 return "anonymous"
1907 session = ctx.session
1908 with _ANON_OWNER_LOCK:
1909 existing = _ANON_OWNER_IDS.get(session)
1910 if existing is None:
1911 existing = f"anon-{uuid.uuid4().hex[:12]}"
1912 _ANON_OWNER_IDS[session] = existing
1913 return existing
1916def _derive_owner(agent_id: str, ctx: Context | None) -> str:
1917 """Resolve the calling agent's stable owner namespace.
1919 Precedence: explicit ``agent_id`` argument, then the ``LILBEE_AGENT_ID`` env var
1920 (pinned in the client's MCP config), then the MCP client name, then a stable
1921 per-connection fallback so unidentified agents never share a namespace.
1922 """
1923 explicit = agent_id or os.environ.get("LILBEE_AGENT_ID", "")
1924 resolved = explicit or _client_name(ctx)
1925 if resolved:
1926 return agent_owner(_slug(resolved))
1927 return agent_owner(_slug(_anon_owner_id(ctx)))
1930@_tool_if(memory_enabled)
1931def memory_remember(
1932 text: str,
1933 kind: MemoryKind = MemoryKind.FACT,
1934 shared: bool = False,
1935 agent_id: str = "",
1936 ctx: Context | None = None,
1937) -> dict[str, Any]:
1938 """Store a durable memory. ``kind``: "fact" (similarity-recalled) or "preference" (always on).
1939 ``shared`` exposes it to the human's TUI/CLI."""
1940 if not memory_enabled():
1941 return _error(MEMORY_DISABLED_HINT)
1942 owner = _derive_owner(agent_id, ctx)
1943 memory_id = remember(text, owner=owner, kind=kind, source=MemorySource.AGENT, shared=shared)
1944 return {"ok": True, "id": memory_id, "owner": owner}
1947@_tool_if(memory_enabled)
1948def memory_recall(
1949 query: str, limit: int = 0, agent_id: str = "", ctx: Context | None = None
1950) -> dict[str, Any]:
1951 """Recall this agent's memories (plus any the human shared) relevant to *query*."""
1952 if not memory_enabled():
1953 return _error(MEMORY_DISABLED_HINT)
1954 owner = _derive_owner(agent_id, ctx)
1955 memories = recall(query, owner, top_k=limit if limit > 0 else None)
1956 return {
1957 "memories": [
1958 {"id": m.id, "text": m.text, "kind": m.kind.value, "owner": m.owner} for m in memories
1959 ]
1960 }
1963@_tool_if(memory_enabled)
1964def memory_list(agent_id: str = "", ctx: Context | None = None) -> dict[str, Any]:
1965 """List every memory in this agent's namespace (any kind, newest first)."""
1966 if not memory_enabled():
1967 return _error(MEMORY_DISABLED_HINT)
1968 owner = _derive_owner(agent_id, ctx)
1969 memories = list_memories(owner)
1970 return {
1971 "memories": [
1972 {"id": m.id, "text": m.text, "kind": m.kind.value, "shared": m.shared} for m in memories
1973 ]
1974 }
1977@_tool_if(memory_enabled)
1978def memory_forget(memory_id: str, agent_id: str = "", ctx: Context | None = None) -> dict[str, Any]:
1979 """Delete one of this agent's own memories by id (agent_id scopes the namespace)."""
1980 if not memory_enabled():
1981 return _error(MEMORY_DISABLED_HINT)
1982 owner = _derive_owner(agent_id, ctx)
1983 if not forget(memory_id, owner=owner):
1984 return _error(f"No memory '{memory_id}' in this agent's namespace.")
1985 return {"ok": True, "id": memory_id}
1988def _placement_dict(view: PlacementView) -> dict[str, Any]:
1989 from lilbee.server.models import PlacementResponse
1991 return PlacementResponse.from_view(view).model_dump(mode="json")
1994def _placement_guard(serialize: Callable[[], dict[str, Any]]) -> dict[str, Any]:
1995 """Run a placement query and serialize it, returning a structured error on failure."""
1996 from lilbee.providers.base import ProviderError
1997 from lilbee.providers.fleet.placement_spec import PlacementError
1999 try:
2000 return serialize()
2001 except (PlacementError, ProviderError) as exc:
2002 return _error(str(exc))
2005def _placement_result(action: Callable[[], PlacementView]) -> dict[str, Any]:
2006 """Run a placement action and serialize its view, returning a structured error on failure."""
2007 return _placement_guard(lambda: _placement_dict(action()))
2010def _parse_spec(spec: dict[str, Any] | None) -> PlacementSpec | None:
2011 from lilbee.providers.fleet.placement_spec import PlacementSpec
2013 return PlacementSpec.from_json(json.dumps(spec)) if spec else None
2016@_tool_named("get_gpus")
2017def get_gpus_tool() -> dict[str, Any]:
2018 """List detected GPUs with free/total VRAM (the placement HTTP /api/gpus equivalent)."""
2020 def _body() -> dict[str, Any]:
2021 from lilbee.cli.tui import messages as msg
2022 from lilbee.providers.fleet.gpu_stats import probe_intel_util_hint
2024 view = get_placement()
2025 hint = probe_intel_util_hint(view.gpus)
2026 return {
2027 "gpus": _placement_dict(view)["gpus"],
2028 "notice": msg.intel_util_hint_text(hint) if hint else None,
2029 }
2031 return _placement_guard(_body)
2034@_tool_named("get_placement")
2035def get_placement_tool() -> dict[str, Any]:
2036 """Show the current effective multi-GPU model placement."""
2037 return _placement_result(get_placement)
2040@_tool_named("preview_placement")
2041def preview_placement_tool(spec: dict[str, Any] | None = None) -> dict[str, Any]:
2042 """Preview what a placement spec (or auto, when omitted) would place. No changes made."""
2043 return _placement_result(lambda: preview_placement(_parse_spec(spec)))
2046@_tool_named("set_placement")
2047def set_placement_tool(spec: dict[str, Any]) -> dict[str, Any]:
2048 """Set and apply a manual multi-GPU placement spec (persists to config).
2050 Maps a role ("chat"/"embed"/"rerank"/"vision") to a placement, e.g.
2051 ``{"chat": {"devices": [0, 1], "tensor_split": [1, 1]}}``. ``devices`` is the
2052 GPU indices (get_gpus lists them); ``tensor_split`` is optional per-device
2053 weights (omit for an even split). A role left out is auto-placed.
2054 """
2055 from lilbee.providers.fleet.placement_spec import PlacementSpec
2057 # set_placement restarts the shared fleet's moved roles: gate it on the
2058 # shared HTTP transport exactly like the REST PUT/DELETE placement routes.
2059 if _transport.http_mounted and not cfg.allow_http_placement:
2060 return _error(placement_refused_message())
2061 # Always build a spec (even {}) so an empty/invalid one is rejected, not cleared.
2062 return _placement_result(lambda: set_placement(PlacementSpec.from_json(json.dumps(spec))))
2065@_tool_named("clear_placement")
2066def clear_placement_tool() -> dict[str, Any]:
2067 """Clear the manual placement and return to automatic placement."""
2068 if _transport.http_mounted and not cfg.allow_http_placement:
2069 return _error(placement_refused_message())
2070 return _placement_result(lambda: set_placement(None))
2073def build_mcp_server() -> LilbeeMCP:
2074 """Build an MCP server carrying every tool registered in this module.
2076 Each transport builds its own instance: the SDK caches one
2077 ``StreamableHTTPSessionManager`` per server and its ``run()`` is single-use,
2078 so a shared server cannot back two apps in one process. Gates registered
2079 via ``_tool_if`` are evaluated here, against current config.
2080 """
2081 server = LilbeeMCP("lilbee", instructions=_INSTRUCTIONS)
2082 for fn, name, gate in _REGISTRATIONS:
2083 if gate is None or gate():
2084 server.add_tool(fn, name=name)
2085 return server
2088_PARENT_DEATH_CLEANUP_S = 5.0
2091def _exit_on_parent_death() -> None:
2092 """Release engine membership best-effort, then hard-exit promptly.
2094 ``os._exit`` skips atexit, so an explicit release stops this process's engine
2095 when it was the last user and keeps the machine clean. But this watchdog's one
2096 contract is to exit promptly on parent death, so the release runs on a daemon
2097 thread joined with a short deadline: a peer holding an engine build lock can
2098 never keep the orphaned process alive with its models resident. The kernel
2099 releases this process's user lock on exit regardless, so a skipped release only
2100 defers the engine stop to the peers' reap and the idle TTL.
2101 """
2102 cleanup = threading.Thread(target=reset_services, name="parent-death-cleanup", daemon=True)
2103 cleanup.start()
2104 cleanup.join(timeout=_PARENT_DEATH_CLEANUP_S)
2105 os._exit(0)
2108async def _serve_stdio(server: LilbeeMCP) -> None:
2109 """Serve over stdio until stdin closes or a signal ends the process."""
2110 exit_on_signals_in_loop(asyncio.get_running_loop())
2111 await server.run_stdio_async()
2114def main() -> None:
2115 """Entry point for the stdio MCP server."""
2116 # Preload so the first tool call doesn't pay the cold-start cost
2117 # of provider/embedder/store init. Failures (missing model, bad
2118 # config) still surface on the first tool call rather than crashing
2119 # the server before it attaches to stdio.
2120 try:
2121 get_services()
2122 except Exception:
2123 log.debug("MCP pre-warm failed; services will init on first call", exc_info=True)
2125 from lilbee.parent_monitor import parse_parent_pid, watch_parent_thread
2127 parent_pid = parse_parent_pid()
2128 if parent_pid is not None:
2129 watch_parent_thread(parent_pid, _exit_on_parent_death)
2131 server = build_mcp_server()
2132 server.follow_config_file()
2133 asyncio.run(_serve_stdio(server))