agentengine-sdk-python 0.6.9__py3-none-any.whl → 0.7.0__py3-none-any.whl

ksadk/sessions/resilient.py

@@ -0,0 +1,360 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import Any, Optional, cast
+
+from ksadk.sessions.base import BaseSessionService, Session, SessionEvent, SessionState
+from ksadk.sessions.in_memory import InMemorySessionService
+from ksadk.sessions.resilience import is_session_backend_failure
+
+logger = logging.getLogger(__name__)
+
+
+class ResilientSessionService(BaseSessionService):
+    """Keep live agent sessions available when durable persistence is unavailable.
+
+    The in-memory service is authoritative for the lifetime of this process. The
+    configured durable service is used as a read-through source and a best-effort
+    write-through sink. After its first failure it stays disabled until a
+    background probe confirms the durable backend is reachable again, at which
+    point it is re-enabled and an INFO log is emitted.
+    """
+
+    _probe_interval_seconds: float = 30.0
+
+    def __init__(
+        self,
+        primary: BaseSessionService,
+        fallback: InMemorySessionService | None = None,
+    ) -> None:
+        self.primary = primary
+        self.fallback = fallback or InMemorySessionService()
+        self._primary_enabled = True
+        self._hydrate_lock = asyncio.Lock()
+        self._primary_session_lock = asyncio.Lock()
+        self._primary_session_ids: set[str] = set()
+        self._probe_task: asyncio.Task[None] | None = None
+
+    @property
+    def degraded(self) -> bool:
+        return not self._primary_enabled
+
+    async def _call_primary(self, method_name: str, *args: Any, **kwargs: Any) -> tuple[bool, Any]:
+        if not self._primary_enabled:
+            return False, None
+        try:
+            method = getattr(self.primary, method_name)
+            return True, await method(*args, **kwargs)
+        except Exception as exc:
+            if not is_session_backend_failure(exc):
+                raise
+            self._disable_primary(exc)
+            return False, None
+
+    def _disable_primary(self, exc: Exception) -> None:
+        if not self._primary_enabled:
+            return
+        self._primary_enabled = False
+        logger.error(
+            "KSADK session persistence degraded; using in-memory live session: %s",
+            exc,
+            extra={
+                "session_backend_state": "degraded",
+                "session_backend": type(self.primary).__name__,
+            },
+        )
+        self._start_probe()
+
+    def _start_probe(self) -> None:
+        if self._probe_task is not None and not self._probe_task.done():
+            return
+        self._probe_task = asyncio.create_task(self._probe_loop())
+
+    async def _probe_loop(self) -> None:
+        while not self._primary_enabled:
+            await asyncio.sleep(self._probe_interval_seconds)
+            if self._primary_enabled:
+                break
+            try:
+                await self.primary.get_session("__ksadk_probe__")
+            except Exception:
+                continue
+            self._primary_enabled = True
+            logger.info(
+                "KSADK session persistence recovered; durable backend re-enabled",
+                extra={
+                    "session_backend_state": "recovered",
+                    "session_backend": type(self.primary).__name__,
+                },
+            )
+
+    async def _hydrate(self, session: Session) -> Session:
+        async with self._hydrate_lock:
+            self._primary_session_ids.add(session.id)
+            existing = await self.fallback.get_session(session.id)
+            if existing is None:
+                await self.fallback.create_session(
+                    session.agent_id,
+                    session.user_id,
+                    session_id=session.id,
+                )
+            existing_events = await self.fallback.get_events(session.id)
+            existing_ids = {event.id for event in existing_events}
+            for event in sorted(session.events, key=lambda item: item.seq_id):
+                if event.id not in existing_ids:
+                    await self.fallback.append_event(session.id, event)
+            await self.fallback.update_session_metadata(
+                session.id,
+                title=session.title,
+                title_source=session.title_source,
+                summary=session.summary,
+                first_prompt=session.first_prompt,
+                last_prompt=session.last_prompt,
+            )
+            current = await self.fallback.get_session(session.id)
+            if current is not None and session.state != current.state:
+                await self.fallback.update_state(
+                    agent_id=session.agent_id,
+                    user_id=session.user_id,
+                    session_id=session.id,
+                    scope="session",
+                    state_delta=session.state,
+                )
+            hydrated = await self.fallback.get_session(session.id)
+            if hydrated is None:
+                raise RuntimeError(f"Failed to hydrate live session {session.id}")
+            return hydrated
+
+    async def create_session(
+        self,
+        agent_id: str,
+        user_id: str,
+        session_id: Optional[str] = None,
+    ) -> Session:
+        if session_id:
+            ok, durable = await self._call_primary("get_session", session_id)
+            if ok and durable is not None:
+                return await self._hydrate(durable)
+            existing = await self.fallback.get_session(session_id)
+            if existing is not None:
+                return existing
+
+        live = await self.fallback.create_session(agent_id, user_id, session_id=session_id)
+        ok, durable = await self._call_primary(
+            "create_session",
+            agent_id,
+            user_id,
+            session_id=live.id,
+        )
+        if ok and durable is not None:
+            self._primary_session_ids.add(durable.id)
+            return await self._hydrate(durable)
+        return live
+
+    async def get_session(self, session_id: str) -> Optional[Session]:
+        live = await self.fallback.get_session(session_id)
+        ok, durable = await self._call_primary("get_session", session_id)
+        if ok and durable is not None:
+            return await self._hydrate(durable)
+        return live
+
+    async def list_sessions(
+        self,
+        agent_id: str,
+        user_id: Optional[str] = None,
+        offset: Optional[int] = None,
+        limit: Optional[int] = None,
+    ) -> list[Session]:
+        ok, durable_sessions = await self._call_primary(
+            "list_sessions",
+            agent_id,
+            user_id,
+            offset,
+            limit,
+        )
+        if ok:
+            for session in durable_sessions or []:
+                await self._hydrate(session)
+        return cast(
+            list[Session],
+            await self.fallback.list_sessions(agent_id, user_id, offset, limit),
+        )
+
+    async def count_sessions(self, agent_id: str, user_id: Optional[str] = None) -> int:
+        sessions = await self.list_sessions(agent_id, user_id)
+        return len(sessions)
+
+    async def delete_session(self, session_id: str) -> bool:
+        deleted = await self.fallback.delete_session(session_id)
+        ok, durable_deleted = await self._call_primary("delete_session", session_id)
+        if ok:
+            self._primary_session_ids.discard(session_id)
+        return deleted or bool(durable_deleted) if ok else deleted
+
+    async def update_session_metadata(
+        self,
+        session_id: str,
+        *,
+        title: Optional[str] = None,
+        title_source: Optional[str] = None,
+        summary: Optional[str] = None,
+        first_prompt: Optional[str] = None,
+        last_prompt: Optional[str] = None,
+    ) -> Session:
+        if await self.fallback.get_session(session_id) is None:
+            await self.get_session(session_id)
+        live = await self.fallback.update_session_metadata(
+            session_id,
+            title=title,
+            title_source=title_source,
+            summary=summary,
+            first_prompt=first_prompt,
+            last_prompt=last_prompt,
+        )
+        await self._ensure_primary_session(session_id)
+        await self._call_primary(
+            "update_session_metadata",
+            session_id,
+            title=title,
+            title_source=title_source,
+            summary=summary,
+            first_prompt=first_prompt,
+            last_prompt=last_prompt,
+        )
+        return live
+
+    async def _ensure_primary_session(self, session_id: str) -> None:
+        """Create the session in PG if it only exists in memory (degraded-era)."""
+        if not self._primary_enabled or session_id in self._primary_session_ids:
+            return
+        async with self._primary_session_lock:
+            if not self._primary_enabled or session_id in self._primary_session_ids:
+                return
+            ok, durable = await self._call_primary("get_session", session_id)
+            if not ok:
+                return
+            if durable is None:
+                live = await self.fallback.get_session(session_id)
+                if live is None:
+                    return
+                ok, durable = await self._call_primary(
+                    "create_session",
+                    live.agent_id,
+                    live.user_id,
+                    session_id=live.id,
+                )
+                if not ok or durable is None:
+                    return
+            self._primary_session_ids.add(session_id)
+
+    async def append_event(self, session_id: str, event: SessionEvent) -> SessionEvent:
+        if await self.fallback.get_session(session_id) is None:
+            await self.get_session(session_id)
+        live = await self.fallback.append_event(session_id, event)
+        await self._ensure_primary_session(session_id)
+        await self._call_primary("append_event", session_id, event)
+        return live
+
+    async def get_events(
+        self,
+        session_id: str,
+        offset: Optional[int] = None,
+        limit: Optional[int] = None,
+        after_seq_id: Optional[int] = None,
+        before_seq_id: Optional[int] = None,
+    ) -> list[SessionEvent]:
+        await self.get_session(session_id)
+        return cast(
+            list[SessionEvent],
+            await self.fallback.get_events(
+                session_id,
+                offset,
+                limit,
+                after_seq_id,
+                before_seq_id,
+            ),
+        )
+
+    async def count_events(
+        self,
+        session_id: str,
+        after_seq_id: Optional[int] = None,
+        before_seq_id: Optional[int] = None,
+    ) -> int:
+        await self.get_session(session_id)
+        return cast(
+            int,
+            await self.fallback.count_events(session_id, after_seq_id, before_seq_id),
+        )
+
+    async def get_state(
+        self,
+        agent_id: str,
+        user_id: Optional[str],
+        session_id: Optional[str],
+        scope: str = "session",
+    ) -> Optional[SessionState]:
+        if session_id:
+            await self.get_session(session_id)
+        live = await self.fallback.get_state(agent_id, user_id, session_id, scope)
+        if live is not None or not self._primary_enabled:
+            return live
+        ok, durable = await self._call_primary(
+            "get_state",
+            agent_id,
+            user_id,
+            session_id,
+            scope,
+        )
+        if ok and durable is not None:
+            return await self.fallback.update_state(
+                agent_id=agent_id,
+                user_id=user_id,
+                session_id=session_id,
+                scope=scope,
+                state_delta=durable.state,
+            )
+        return live
+
+    async def update_state(
+        self,
+        *,
+        agent_id: str,
+        user_id: Optional[str],
+        session_id: Optional[str],
+        scope: str,
+        state_delta: dict[str, Any],
+    ) -> SessionState:
+        if session_id and await self.fallback.get_session(session_id) is None:
+            await self.get_session(session_id)
+        live = await self.fallback.update_state(
+            agent_id=agent_id,
+            user_id=user_id,
+            session_id=session_id,
+            scope=scope,
+            state_delta=state_delta,
+        )
+        if session_id:
+            await self._ensure_primary_session(session_id)
+        await self._call_primary(
+            "update_state",
+            agent_id=agent_id,
+            user_id=user_id,
+            session_id=session_id,
+            scope=scope,
+            state_delta=state_delta,
+        )
+        return live
+
+    async def aclose(self) -> None:
+        if self._probe_task is not None and not self._probe_task.done():
+            self._probe_task.cancel()
+            try:
+                await self._probe_task
+            except asyncio.CancelledError:
+                pass
+        for service in (self.primary, self.fallback):
+            close = getattr(service, "aclose", None)
+            if close is not None:
+                await close()

ksadk/toolsets/__init__.py

@@ -53,6 +53,8 @@ _FOCUSED_TOOL_NAMES = (
     "workspace_status",
     "list_workspace_files",
     "read_workspace_file",
+    "write_workspace_file",
+    "write_workspace_files",
     "search_workspace_files",
     "edit_workspace_file",
     "multi_edit_workspace_file",

ksadk/toolsets/workspace.py

@@ -150,7 +150,10 @@ def read_workspace_file(
     max_chars: int | None = None,
     include_line_numbers: bool = True,
 ) -> dict[str, Any]:
-    """Read a UTF-8 text file from the AgentEngine workspace."""
+    """Read a UTF-8 workspace file. Inspect the result before editing it.
+
+    Do not issue a dependent read and edit in the same parallel tool-call batch.
+    """
 
     return _gateway().invoke(
         "read_workspace_file",
@@ -313,7 +316,10 @@ def edit_workspace_file(
     replace_all: bool = False,
     approval: dict[str, Any] | None = None,
 ) -> dict[str, Any]:
-    """Replace an exact text snippet inside a UTF-8 workspace file."""
+    """Replace exact text after reading the file in a prior model step.
+
+    The read must complete before this call; do not batch a dependent read and edit.
+    """
 
     return _gateway().invoke(
         "edit_workspace_file",
@@ -336,6 +342,9 @@ def _edit_workspace_file_impl(
 ) -> dict[str, Any]:
     target = resolve_workspace_path(path)
     relative = workspace_relative(target)
+    request_error = _validate_single_edit_request(old_text, relative=relative)
+    if request_error is not None:
+        return request_error
     read_error = _validate_read_state(target, relative)
     if read_error is not None:
         return read_error
@@ -376,7 +385,10 @@ def multi_edit_workspace_file(
     replace_all: bool = False,
     approval: dict[str, Any] | None = None,
 ) -> dict[str, Any]:
-    """Apply multiple exact snippet edits atomically inside one workspace file."""
+    """Apply atomic exact edits after reading the file in a prior model step.
+
+    The read must complete before this call; do not batch a dependent read and edit.
+    """
 
     return _gateway().invoke(
         "multi_edit_workspace_file",
@@ -395,21 +407,36 @@ def _multi_edit_workspace_file_impl(
 ) -> dict[str, Any]:
     target = resolve_workspace_path(path)
     relative = workspace_relative(target)
+    if not isinstance(edits, list) or not edits:
+        return {
+            "ok": False,
+            "path": relative,
+            "error_type": "invalid_edits",
+            "error_message": "edits must be a non-empty list",
+        }
+    for index, edit in enumerate(edits):
+        if not isinstance(edit, dict):
+            return {
+                "ok": False,
+                "path": relative,
+                "error_type": "invalid_edit",
+                "error_message": f"edit at index {index} must be an object",
+                "failed_edit_index": index,
+            }
+        request_error = _validate_single_edit_request(edit.get("old_text"), relative=relative)
+        if request_error is not None:
+            return {**request_error, "failed_edit_index": index, "edit_count": len(edits)}
     read_error = _validate_read_state(target, relative)
     if read_error is not None:
         return read_error
     text, error = _read_workspace_text(target)
     if error is not None:
         return {**error, "path": path}
-    if not isinstance(edits, list) or not edits:
-        return {"ok": False, "path": relative, "error_type": "invalid_edits", "error_message": "edits must be a non-empty list"}
     original = text or ""
     updated = original
     diagnostics: list[dict[str, Any]] = []
     total_replacements = 0
     for index, edit in enumerate(edits):
-        if not isinstance(edit, dict):
-            return {"ok": False, "path": relative, "error_type": "invalid_edit", "error_message": f"edit at index {index} must be an object", "failed_edit_index": index}
         edit_result = _apply_single_edit(
             updated,
             edit.get("old_text"),
@@ -718,9 +745,10 @@ def _apply_single_edit(
     replace_all: bool = False,
     relative: str,
 ) -> dict[str, Any]:
-    old = str(old_text or "")
-    if not old:
-        return {"ok": False, "path": relative, "error_type": "missing_old_text", "error_message": "old_text is required"}
+    request_error = _validate_single_edit_request(old_text, relative=relative)
+    if request_error is not None:
+        return request_error
+    old = str(old_text)
     expected = max(1, int(expected_replacements or 1))
     actual_old, used_quote_normalization = _find_actual_string(text, old)
     matches = _match_previews(text, actual_old or old) if actual_old else []
@@ -756,6 +784,17 @@ def _apply_single_edit(
     }
 
 
+def _validate_single_edit_request(old_text: Any, *, relative: str) -> dict[str, Any] | None:
+    if not str(old_text or ""):
+        return {
+            "ok": False,
+            "path": relative,
+            "error_type": "missing_old_text",
+            "error_message": "old_text is required",
+        }
+    return None
+
+
 def _match_previews(text: str, needle: str, limit: int = 20) -> list[dict[str, Any]]:
     if not needle:
         return []

ksadk/tracing/setup.py

@@ -174,6 +174,16 @@ def _has_nonzero_token_usage(usage: dict[str, int]) -> bool:
     return any(value > 0 for value in usage.values())
 
 
+def _instrumentation_scope_name(span: Any) -> str:
+    scope = getattr(span, "instrumentation_scope", None)
+    name = getattr(scope, "name", None)
+    if name:
+        return str(name)
+    info = getattr(span, "instrumentation_info", None)
+    name = getattr(info, "name", None)
+    return str(name or "")
+
+
 def _clone_span_with_attributes(span: Any, attributes: dict[str, Any]) -> Any:
     try:
         from opentelemetry.sdk.trace import ReadableSpan
@@ -198,6 +208,53 @@ def _clone_span_with_attributes(span: Any, attributes: dict[str, Any]) -> Any:
         return span
 
 
+def _prepare_langfuse_spans(spans: Any) -> Any:
+    """Keep ksadk response usage authoritative for Langfuse exports.
+
+    OpenInference LangChain instrumentation can emit llm.token_count.* on nested
+    ChatOpenAI spans. In LangGraph streaming these counts can be cumulative and
+    disagree with the usage returned by ksadk to RunAgent/session state. When a
+    trace already has a ksadk.conversations span with gen_ai.usage.*, strip token
+    attrs from nested OpenInference LangChain spans before exporting to
+    Langfuse. The child spans remain useful for timing and structure, but cannot
+    pollute token accounting.
+    """
+    if not spans:
+        return spans
+    span_list = list(spans)
+    authoritative_trace_ids = {
+        _trace_id(span)
+        for span in span_list
+        if _trace_id(span) is not None
+        and _instrumentation_scope_name(span) == "ksadk.conversations"
+        and _has_nonzero_token_usage(_span_token_usage(span))
+    }
+    if not authoritative_trace_ids:
+        return spans
+
+    strip_prefixes = ("llm.token_count.", "llm.usage.", "gen_ai.usage.")
+    replacements: dict[int, Any] = {}
+    for span in span_list:
+        if _trace_id(span) not in authoritative_trace_ids:
+            continue
+        if _instrumentation_scope_name(span) != "openinference.instrumentation.langchain":
+            continue
+        attributes = dict(getattr(span, "attributes", None) or {})
+        stripped = {
+            key: value
+            for key, value in attributes.items()
+            if not any(str(key).startswith(prefix) for prefix in strip_prefixes)
+        }
+        if len(stripped) != len(attributes):
+            if str(stripped.get("openinference.span.kind") or "").upper() == "LLM":
+                stripped["openinference.span.kind"] = "CHAIN"
+            replacements[id(span)] = _clone_span_with_attributes(span, stripped)
+
+    if not replacements:
+        return spans
+    return [replacements.get(id(span), span) for span in span_list]
+
+
 def _prepare_cloud_monitor_spans(spans: Any) -> Any:
     """Add root token rollups for CloudMonitor without changing other exporters.
 
@@ -634,6 +691,11 @@ def setup_tracing(
                 endpoint=generic_otlp_config["endpoint"],
                 service_name=_get_service_name(),
                 header_keys=sorted(generic_otlp_config["headers"]),
+                span_transform=(
+                    _prepare_langfuse_spans
+                    if _is_langfuse_otlp_endpoint(generic_otlp_config["endpoint"])
+                    else None
+                ),
             )
             provider.add_span_processor(
                 BatchSpanProcessor(
@@ -721,6 +783,7 @@ def setup_tracing(
                     endpoint=config["endpoint"],
                     service_name=_get_service_name(),
                     header_keys=sorted(config["headers"]),
+                    span_transform=_prepare_langfuse_spans,
                 )
                 provider.add_span_processor(
                 BatchSpanProcessor(

ksadk/tui/__init__.py

@@ -1,12 +1,27 @@
 """
-KsADK TUI - Textual Terminal User Interface
+KsADK TUI - prompt_toolkit 全屏交互（对标 Codex CLI）
 
-提供简洁的交互体验:
-- 流式 Markdown 输出
-- 历史记录支持
-- 思考过程显示
+- 全屏 alternate screen，transcript 区（FormattedTextControl + ANSI 保留 rich 颜色）
+  + 底部输入框 + footer 状态栏
+- 滚动：PageUp/PgDn/↑↓（cursor 跟随视口顶 + 手动 vertical_scroll）
+- 流式：streaming 期动态区累积，turn 结束落定整条 rich 渲染（历史 ANSI 缓存）
+- 命令：/new /clear /session /model（模型热切 picker）? exit；Ctrl-C 取消流式/退出
 """
 
-from ksadk.tui.app import AgentTUI
+from ksadk.tui.loop import (
+    InteractionLoop,
+    InterruptPending,
+    RichLiveRenderer,
+    TranscriptEntry,
+    render_stream,
+    run_tui,
+)
 
-__all__ = ["AgentTUI"]
+__all__ = [
+    "InteractionLoop",
+    "InterruptPending",
+    "RichLiveRenderer",
+    "TranscriptEntry",
+    "render_stream",
+    "run_tui",
+]