kash-shell 0.3.30__py3-none-any.whl → 0.3.33__py3-none-any.whl

kash/file_storage/item_id_index.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+import threading
+from collections.abc import Callable
+
+from prettyfmt import fmt_lines, fmt_path
+
+from kash.config.logger import get_logger
+from kash.file_storage.store_filenames import join_suffix, parse_item_filename
+from kash.model.items_model import Item, ItemId
+from kash.model.paths_model import StorePath
+from kash.utils.common.uniquifier import Uniquifier
+from kash.utils.errors import InvalidFilename, SkippableError
+
+log = get_logger(__name__)
+
+
+class ItemIdIndex:
+    """
+    Index of item identities and historical filenames within a workspace.
+
+    - Tracks a mapping of `ItemId -> StorePath` for quick lookups
+    - Tracks historical slugs via `Uniquifier` to generate unique names consistently
+
+    TODO: Should add a file system watcher to make this always consistent with disk state.
+    """
+
+    def __init__(self) -> None:
+        self._lock = threading.RLock()
+        self.uniquifier = Uniquifier()
+        self.id_map: dict[ItemId, StorePath] = {}
+
+    def reset(self) -> None:
+        """
+        Clear all index state.
+        """
+        with self._lock:
+            log.info("ItemIdIndex: reset")
+            self.uniquifier = Uniquifier()
+            self.id_map.clear()
+
+    def __len__(self) -> int:
+        """
+        Number of unique names tracked.
+        """
+        with self._lock:
+            return len(self.uniquifier)
+
+    def uniquify_slug(self, slug: str, full_suffix: str) -> tuple[str, list[str]]:
+        """
+        Return a unique slug and historic slugs for the given suffix.
+        """
+        with self._lock:
+            # This updates internal history as a side-effect. Log for consistency.
+            log.info("ItemIdIndex: uniquify slug '%s' with suffix '%s'", slug, full_suffix)
+            return self.uniquifier.uniquify_historic(slug, full_suffix)
+
+    def index_item(
+        self, store_path: StorePath, load_item: Callable[[StorePath], Item]
+    ) -> StorePath | None:
+        """
+        Update the index with an item at `store_path`.
+        Returns store path of any duplicate item with the same id, otherwise None.
+        """
+        name, item_type, _format, file_ext = parse_item_filename(store_path)
+        if not file_ext:
+            log.debug(
+                "Skipping file with unrecognized name or extension: %s",
+                fmt_path(store_path),
+            )
+            return None
+
+        with self._lock:
+            full_suffix = join_suffix(item_type.name, file_ext.name) if item_type else file_ext.name
+            # Track unique name history
+            self.uniquifier.add(name, full_suffix)
+
+        log.info("ItemIdIndex: indexing %s", fmt_path(store_path))
+
+        # Load item outside the lock to avoid holding it during potentially slow I/O
+        try:
+            item = load_item(store_path)
+        except (ValueError, SkippableError) as e:
+            log.warning(
+                "ItemIdIndex: could not index file, skipping: %s: %s",
+                fmt_path(store_path),
+                e,
+            )
+            return None
+
+        dup_path: StorePath | None = None
+        with self._lock:
+            item_id = item.item_id()
+            if item_id:
+                old_path = self.id_map.get(item_id)
+                if old_path and old_path != store_path:
+                    dup_path = old_path
+                    log.info(
+                        "ItemIdIndex: duplicate id detected %s:\n%s",
+                        item_id,
+                        fmt_lines([old_path, store_path]),
+                    )
+                self.id_map[item_id] = store_path
+                log.info("ItemIdIndex: set id %s -> %s", item_id, fmt_path(store_path))
+
+        return dup_path
+
+    def unindex_item(self, store_path: StorePath, load_item: Callable[[StorePath], Item]) -> None:
+        """
+        Remove an item from the id index.
+        """
+        try:
+            # Load item outside the lock to avoid holding it during potentially slow I/O
+            item = load_item(store_path)
+            item_id = item.item_id()
+            if item_id:
+                with self._lock:
+                    try:
+                        self.id_map.pop(item_id, None)
+                        log.info("ItemIdIndex: removed id %s for %s", item_id, fmt_path(store_path))
+                    except KeyError:
+                        pass
+        except (FileNotFoundError, InvalidFilename):
+            pass
+
+    def find_store_path_by_id(self, item_id: ItemId) -> StorePath | None:
+        with self._lock:
+            return self.id_map.get(item_id)

kash/mcp/mcp_server_routes.py

@@ -6,8 +6,10 @@ from dataclasses import dataclass
 
 from funlog import log_calls
 from mcp.server.lowlevel import Server
+from mcp.server.lowlevel.server import StructuredContent, UnstructuredContent
 from mcp.types import Prompt, Resource, TextContent, Tool
 from prettyfmt import fmt_path
+from pydantic import BaseModel
 from strif import AtomicVar
 
 from kash.config.capture_output import CapturedOutput, captured_output
@@ -20,6 +22,7 @@ from kash.model.actions_model import Action, ActionResult
 from kash.model.exec_model import ExecContext
 from kash.model.params_model import TypedParamValues
 from kash.model.paths_model import StorePath
+from kash.utils.common.url import Url
 
 log = get_logger(__name__)
 
@@ -109,6 +112,22 @@ def get_published_tools() -> list[Tool]:
         return []
 
 
+class StructuredActionResult(BaseModel):
+    """
+    Error from an MCP tool call.
+    """
+
+    s3_paths: list[Url] | None = None
+    """If the tool created an S3 item, the S3 paths of the created items."""
+
+    error: str | None = None
+    """If the tool had an error, the error message."""
+
+    # TODO: Include other metadata.
+    # metadata: dict[str, Any] | None = None
+    # """Metadata about the action result."""
+
+
 @dataclass(frozen=True)
 class ToolResult:
     """
@@ -119,6 +138,7 @@ class ToolResult:
     captured_output: CapturedOutput
     action_result: ActionResult
     result_store_paths: list[StorePath]
+    result_s3_paths: list[Url]
     error: Exception | None = None
 
     @property
@@ -168,12 +188,13 @@ class ToolResult:
         # TODO: Add more info on how to find the logs.
         return "Check kash logs for details."
 
-    def formatted_for_client(self) -> list[TextContent]:
+    def as_mcp_content(self) -> tuple[UnstructuredContent, StructuredContent]:
         """
-        Convert the tool result to content for the client LLM.
+        Convert the tool result to content for the MCP client.
         """
+        structured = StructuredActionResult()
         if self.error:
-            return [
+            unstructured = [
                 TextContent(
                     text=f"The tool `{self.action.name}` had an error: {self.error}.\n\n"
                     + self.check_logs_message,
@@ -194,7 +215,7 @@ class ToolResult:
             if not chat_result:
                 chat_result = "No result. Check kash logs for details."
 
-            return [
+            unstructured = [
                 TextContent(
                     text=f"{self.output_summary}\n\n"
                     f"{self.output_content}\n\n"
@@ -202,10 +223,15 @@ class ToolResult:
                     type="text",
                 ),
             ]
+            structured = StructuredActionResult(s3_paths=self.result_s3_paths)
+
+        return unstructured, structured.model_dump()
 
 
 @log_calls(level="info")
-def run_mcp_tool(action_name: str, arguments: dict) -> list[TextContent]:
+def run_mcp_tool(
+    action_name: str, arguments: dict
+) -> tuple[UnstructuredContent, StructuredContent]:
     """
     Run the action as a tool.
     """
@@ -222,6 +248,7 @@ def run_mcp_tool(action_name: str, arguments: dict) -> list[TextContent]:
                 refetch=False,  # Using the file caches.
                 # Keeping all transient files for now, but maybe make transient?
                 override_state=None,
+                sync_to_s3=True,  # Enable S3 syncing for MCP tools.
             ) as exec_settings:
                 action_cls = look_up_action_class(action_name)
 
@@ -237,9 +264,9 @@ def run_mcp_tool(action_name: str, arguments: dict) -> list[TextContent]:
                 context = ExecContext(action=action, settings=exec_settings)
                 action_input = prepare_action_input(*input_items)
 
-                result, result_store_paths, _archived_store_paths = run_action_with_caching(
-                    context, action_input
-                )
+                result_with_paths = run_action_with_caching(context, action_input)
+                result = result_with_paths.result
+                result_store_paths = result_with_paths.result_paths
 
         # Return final result, formatted for the LLM to understand.
         return ToolResult(
@@ -247,8 +274,9 @@ def run_mcp_tool(action_name: str, arguments: dict) -> list[TextContent]:
             captured_output=capture.output,
             action_result=result,
             result_store_paths=result_store_paths,
+            result_s3_paths=result_with_paths.s3_paths,
             error=None,
-        ).formatted_for_client()
+        ).as_mcp_content()
 
     except Exception as e:
         log.exception("Error running mcp tool")
@@ -258,7 +286,7 @@ def run_mcp_tool(action_name: str, arguments: dict) -> list[TextContent]:
                 + "Check kash logs for details.",
                 type="text",
             )
-        ]
+        ], StructuredActionResult(error=str(e)).model_dump()
 
 
 def create_base_server() -> Server:
@@ -288,7 +316,9 @@ def create_base_server() -> Server:
         return []
 
     @app.call_tool()
-    async def handle_tool(name: str, arguments: dict) -> list[TextContent]:
+    async def handle_tool(
+        name: str, arguments: dict
+    ) -> tuple[UnstructuredContent, StructuredContent]:
         try:
             if name not in _mcp_published_actions.copy():
                 log.error(f"Unknown tool requested: {name}")
@@ -303,6 +333,6 @@ def create_base_server() -> Server:
                     text=f"Error executing tool {name}: {e}",
                     type="text",
                 )
-            ]
+            ], StructuredActionResult(error=str(e)).model_dump()
 
     return app

kash/model/actions_model.py

@@ -246,7 +246,17 @@ class Action(ABC):
 
     output_type: ItemType = ItemType.doc
     """
-    The type of the output item(s), which for now are all assumed to be of the same type.
+    The type of the output item(s). If an action returns multiple output types,
+    this will be the output type of the first output.
+    This is mainly used for preassembly for the cache check if an output already exists.
+    """
+
+    output_format: Format | None = None
+    """
+    The format of the output item(s). The default is to assume it is the same
+    format as the input. If an action returns multiple output formats,
+    this will be the format of the first output.
+    This is mainly used for preassembly for the cache check if an output already exists.
     """
 
     expected_outputs: ArgCount = ONE_ARG
@@ -540,7 +550,7 @@ class Action(ABC):
         """
         can_preassemble = self.cacheable and self.expected_outputs == ONE_ARG
         log.info(
-            "Preassemble check for `%s` is %s (%s with cacheable=%s)",
+            "Preassemble check for `%s`: can_preassemble=%s (expected_outputs=%s, cacheable=%s)",
             self.name,
             can_preassemble,
             self.expected_outputs,
@@ -549,9 +559,10 @@ class Action(ABC):
         if can_preassemble:
             # Using first input to determine the output title.
             primary_input = context.action_input.items[0]
-            # In this case we only expect one output.
-            item = primary_input.derived_copy(context, 0)
-            return ActionResult([item])
+            # In this case we only expect one output, of the type specified by the action.
+            primary_output = primary_input.derived_copy(context, 0, type=context.action.output_type)
+            log.info("Preassembled output: source %s, %s", primary_output.source, primary_output)
+            return ActionResult([primary_output])
         else:
             # Caching disabled.
             return None

kash/model/exec_model.py

@@ -43,6 +43,9 @@ class RuntimeSettings:
     no_format: bool = False
     """If True, will not normalize the output item's body text formatting (for Markdown)."""
 
+    sync_to_s3: bool = True
+    """If True, will sync output items to S3 if input was from S3."""
+
     @property
     def workspace(self) -> FileStore:
         from kash.workspaces.workspaces import get_ws

kash/model/items_model.py

@@ -203,6 +203,15 @@ class ItemId:
             # If we got here, the item has no identity.
             item_id = None
 
+        log.debug(
+            "item_id is %s for type=%s, format=%s, url=%s, title=%s, source=%s",
+            item_id,
+            item.type,
+            item.format,
+            item.url,
+            item.title,
+            item.source,
+        )
         return item_id
 
 
@@ -835,7 +844,9 @@ class Item:
         the type and the body.
 
         Same as `new_copy_with` but also updates the `derived_from` relation. If we also
-        have an action context, then use the `title_template` to derive a new title.
+        have an action context, then use that to fill some fields, in particular `title_template`
+        to derive a new title and `output_type` and `output_format` to set the output type
+        and format
         """
 
         # Get derived_from relation if possible.
@@ -869,20 +880,38 @@ class Item:
         if "external_path" not in updates:
             updates["external_path"] = None
 
+        action_context = action_context or self.context
+
+        if action_context:
+            # Default the output item type and format to the action's declared output_type
+            # and format if not explicitly set.
+            if "type" not in updates:
+                updates["type"] = action_context.action.output_type
+            # If we were not given a format override, we leave the output type the same.
+            elif action_context.action.output_format:
+                # Check an overridden format and then our own format.
+                new_output_format = updates.get("format", self.format)
+                if new_output_format and action_context.action.output_format != new_output_format:
+                    log.warning(
+                        "Output item format `%s` does not match declared output format `%s` for action `%s`",
+                        new_output_format,
+                        action_context.action.output_format,
+                        action_context.action.name,
+                    )
+
         new_item = self.new_copy_with(update_timestamp=True, **updates)
         if derived_from:
             new_item.update_relations(derived_from=derived_from)
 
-        action_context = action_context or self.context
-
         # Record the history.
         if action_context:
-            self.source = Source(
-                operation=action_context.operation,
-                output_num=output_num,
-                cacheable=action_context.action.cacheable,
+            new_item.update_source(
+                Source(
+                    operation=action_context.operation,
+                    output_num=output_num,
+                    cacheable=action_context.action.cacheable,
+                )
             )
-            self.add_to_history(self.source.operation.summary())
             action = action_context.action
         else:
             action = None
@@ -911,9 +940,10 @@ class Item:
             setattr(self.relations, key, list(value))
         return self.relations
 
-    def update_history(self, source: Source) -> None:
+    def update_source(self, source: Source) -> None:
         """
-        Update the history of the item with the given operation.
+        Update the source and the history of the item to indicate it was created
+        by the given operation. For convenience, this is idempotent.
         """
         self.source = source
         self.add_to_history(source.operation.summary())
@@ -945,6 +975,9 @@ class Item:
         return metadata_matches and body_matches
 
     def add_to_history(self, operation_summary: OperationSummary):
+        """
+        For convenience, this is idempotent.
+        """
         if not self.history:
             self.history = []
         # Don't add duplicates to the history.

kash/utils/api_utils/gather_limited.py

@@ -542,6 +542,8 @@ async def gather_limited_sync(
             # Mark as failed
             if status and task_id is not None:
                 await status.finish(task_id, TaskState.FAILED, str(e))
+
+            log.warning("Task failed: %s: %s", label, e, exc_info=True)
             raise
 
     return await _gather_with_interrupt_handling(

kash/utils/api_utils/multitask_gather.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+from collections.abc import Callable, Iterable, Sequence
+from typing import Any, TypeVar
+
+from kash.config.logger import get_logger
+from kash.config.settings import global_settings
+from kash.shell.output.shell_output import multitask_status
+from kash.utils.api_utils.api_retries import RetrySettings
+from kash.utils.api_utils.gather_limited import FuncTask, Limit, gather_limited_sync
+
+T = TypeVar("T")
+
+log = get_logger(name=__name__)
+
+
+def _default_labeler(total: int) -> Callable[[int, Any], str]:
+    def labeler(i: int, _spec: Any) -> str:  # pyright: ignore[reportUnusedParameter]
+        return f"Task {i + 1}/{total}"
+
+    return labeler
+
+
+async def multitask_gather(
+    tasks: Iterable[FuncTask[T]] | Sequence[FuncTask[T]],
+    *,
+    labeler: Callable[[int, Any], str] | None = None,
+    limit: Limit | None = None,
+    bucket_limits: dict[str, Limit] | None = None,
+    retry_settings: RetrySettings | None = None,
+    show_progress: bool = True,
+) -> list[T]:
+    """
+    Run many `FuncTask`s concurrently with shared progress UI and rate limits.
+
+    This wraps the standard pattern of creating a status context, providing a labeler,
+    and calling `gather_limited_sync` with common options.
+
+    - `labeler` can be omitted; a simple "Task X/Y" label will be used.
+    - If `limit` is not provided, defaults are taken from `global_settings()`.
+    - If `show_progress` is False, tasks are run without the status context.
+    - By default, exceptions are returned as results rather than raised (return_exceptions=True).
+    """
+
+    # Normalize tasks to a list for length and stable iteration
+    task_list: list[FuncTask[T]] = list(tasks)
+
+    # Provide a default labeler if none is supplied
+    effective_labeler: Callable[[int, Any], str] = (
+        labeler if labeler is not None else _default_labeler(len(task_list))
+    )
+
+    # Provide sensible default rate limits if none are supplied
+    effective_limit: Limit = (
+        limit
+        if limit is not None
+        else Limit(
+            rps=global_settings().limit_rps,
+            concurrency=global_settings().limit_concurrency,
+        )
+    )
+
+    if not show_progress:
+        log.warning("Running %d tasks (progress disabled)…", len(task_list))
+
+    async with multitask_status(enabled=show_progress) as status:
+        return await gather_limited_sync(
+            *task_list,
+            limit=effective_limit,
+            bucket_limits=bucket_limits,
+            status=status,
+            labeler=effective_labeler,
+            retry_settings=retry_settings,
+        )

kash/utils/common/s3_utils.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+import shutil
+import subprocess
+from pathlib import Path
+
+from sidematter_format.sidematter_format import Sidematter
+
+from kash.utils.common.url import Url, is_s3_url, parse_s3_url
+
+
+def check_aws_cli() -> None:
+    """
+    Check if the AWS CLI is installed and available.
+    """
+    if shutil.which("aws") is None:
+        raise RuntimeError(
+            "AWS CLI not found in PATH. Please install 'awscli' and ensure 'aws' is available."
+        )
+
+
+def get_s3_parent_folder(url: Url) -> Url | None:
+    """
+    Get the parent folder of an S3 URL, or None if not an S3 URL.
+    """
+    if is_s3_url(url):
+        s3_bucket, s3_key = parse_s3_url(url)
+        s3_parent_folder = Path(s3_key).parent
+
+        return Url(f"s3://{s3_bucket}/{s3_parent_folder}")
+
+    else:
+        return None
+
+
+def s3_sync_to_folder(
+    src_path: str | Path,
+    s3_dest_parent: Url,
+    *,
+    include_sidematter: bool = False,
+) -> list[Url]:
+    """
+    Sync a local file or directory to an S3 "parent" folder using the AWS CLI.
+    Set `include_sidematter` to include sidematter files alongside the source files.
+
+    Returns a list of S3 URLs that were the top-level sync targets:
+    - For a single file: the file URL (and sidematter file/dir URLs if included).
+    - For a directory: the destination parent prefix URL (non-recursive reporting).
+    """
+
+    src_path = Path(src_path)
+    if not src_path.exists():
+        raise ValueError(f"Source path does not exist: {src_path}")
+    if not is_s3_url(s3_dest_parent):
+        raise ValueError(f"Destination must be an s3:// URL: {s3_dest_parent}")
+
+    check_aws_cli()
+
+    dest_prefix = str(s3_dest_parent).rstrip("/") + "/"
+    targets: list[Url] = []
+
+    if src_path.is_file():
+        # Build the list of paths to sync using Sidematter's resolved path_list if requested.
+        sync_paths: list[Path]
+        if include_sidematter:
+            resolved = Sidematter(src_path).resolve(parse_meta=False, use_frontmatter=False)
+            sync_paths = resolved.path_list
+        else:
+            sync_paths = [src_path]
+
+        for p in sync_paths:
+            if p.is_file():
+                # Use sync with include/exclude to leverage default short-circuiting
+                subprocess.run(
+                    [
+                        "aws",
+                        "s3",
+                        "sync",
+                        str(p.parent),
+                        dest_prefix,
+                        "--exclude",
+                        "*",
+                        "--include",
+                        p.name,
+                    ],
+                    check=True,
+                )
+                targets.append(Url(dest_prefix + p.name))
+            elif p.is_dir():
+                dest_dir = dest_prefix + p.name + "/"
+                subprocess.run(["aws", "s3", "sync", str(p), dest_dir], check=True)
+                targets.append(Url(dest_dir))
+
+        return targets
+    else:
+        # Directory mode: sync whole directory.
+        subprocess.run(
+            [
+                "aws",
+                "s3",
+                "sync",
+                str(src_path),
+                dest_prefix,
+            ],
+            check=True,
+        )
+        targets.append(Url(dest_prefix))
+        return targets

kash/utils/common/url.py

@@ -26,6 +26,7 @@ A string that may not be resolved to a URL or path.
 
 HTTP_ONLY = ["http", "https"]
 HTTP_OR_FILE = HTTP_ONLY + ["file"]
+HTTP_OR_FILE_OR_S3 = HTTP_OR_FILE + ["s3"]
 
 
 def check_if_url(
@@ -36,7 +37,8 @@ def check_if_url(
     the `urlparse.ParseResult`.
 
     Also returns false for Paths, so that it's easy to use local paths and URLs
-    (`Locator`s) interchangeably. Can provide `HTTP_ONLY` or `HTTP_OR_FILE` to
+    (`Locator`s) interchangeably. Can provide `HTTP_ONLY` or `HTTP_OR_FILE`
+    or `HTTP_OR_FILE_OR_S3` to restrict to only certain schemes.
     restrict to only certain schemes.
     """
     if isinstance(text, Path):
@@ -69,6 +71,13 @@ def is_file_url(url: str | Url) -> bool:
     return url.startswith("file://")
 
 
+def is_s3_url(url: str | Url) -> bool:
+    """
+    Is URL an S3 URL?
+    """
+    return url.startswith("s3://")
+
+
 def parse_http_url(url: str | Url) -> ParseResult:
     """
     Parse an http/https URL and return the parsed result, raising ValueError if
@@ -118,7 +127,7 @@ def as_file_url(path: str | Path) -> Url:
 
 def normalize_url(
     url: Url,
-    check_schemes: list[str] | None = HTTP_OR_FILE,
+    check_schemes: list[str] | None = HTTP_OR_FILE_OR_S3,
     drop_fragment: bool = True,
     resolve_local_paths: bool = True,
 ) -> Url:
@@ -238,7 +247,10 @@ def test_normalize_url():
         normalize_url(url=Url("/not/a/URL"))
         raise AssertionError()
     except ValueError as e:
-        assert str(e) == "Scheme '' not in allowed schemes: ['http', 'https', 'file']: /not/a/URL"
+        assert (
+            str(e)
+            == "Scheme '' not in allowed schemes: ['http', 'https', 'file', 's3']: /not/a/URL"
+        )
 
     try:
         normalize_url(Url("ftp://example.com"))
@@ -246,7 +258,7 @@ def test_normalize_url():
     except ValueError as e:
         assert (
             str(e)
-            == "Scheme 'ftp' not in allowed schemes: ['http', 'https', 'file']: ftp://example.com"
+            == "Scheme 'ftp' not in allowed schemes: ['http', 'https', 'file', 's3']: ftp://example.com"
         )