@openlucaskaka/kagent 0.1.7

package/src/kagent/service/safety.py

@@ -0,0 +1,139 @@
+from __future__ import annotations
+
+import hmac
+from ipaddress import ip_address
+from typing import Any, Mapping
+from uuid import uuid4
+
+_MAX_REQUEST_ID_LENGTH = 128
+_MAX_IDEMPOTENCY_KEY_LENGTH = 128
+_SAFE_TRACE_FILENAME_CHARS = frozenset(
+    "abcdefghijklmnopqrstuvwxyz"
+    "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+    "0123456789"
+    "._-"
+)
+
+
+def authorized(
+    headers: Mapping[str, str],
+    auth_token: str,
+    auth_tokens: Mapping[str, str] | None = None,
+) -> bool:
+    if not auth_token and not auth_tokens:
+        return True
+    return bool(authenticated_subject(headers, auth_token, auth_tokens or {}))
+
+
+def authenticated_subject(
+    headers: Mapping[str, str],
+    auth_token: str,
+    auth_tokens: Mapping[str, str] | None = None,
+) -> str:
+    provided = header_value(headers, "Authorization")
+    if not safe_header_value(provided):
+        return ""
+    if auth_token:
+        expected = f"Bearer {auth_token}"
+        if safe_header_value(expected) and hmac.compare_digest(provided, expected):
+            return "default"
+    for subject, token in sorted((auth_tokens or {}).items()):
+        expected = f"Bearer {token}"
+        if safe_header_value(expected) and hmac.compare_digest(provided, expected):
+            return subject
+    return ""
+
+
+def authenticated_with_primary_token(headers: Mapping[str, str], auth_token: str) -> bool:
+    if not auth_token:
+        return False
+    provided = header_value(headers, "Authorization")
+    expected = f"Bearer {auth_token}"
+    return (
+        safe_header_value(provided)
+        and safe_header_value(expected)
+        and hmac.compare_digest(provided, expected)
+    )
+
+
+def json_content_type(headers: Mapping[str, str]) -> bool:
+    content_type = header_value(headers, "Content-Type")
+    if not content_type:
+        return True
+    media_type = content_type.split(";", 1)[0].strip().lower()
+    return media_type == "application/json"
+
+
+def rate_limit_key(
+    headers: Mapping[str, str],
+    remote_addr: str,
+    *,
+    trust_forwarded_for: bool,
+    auth_token: str = "",
+    auth_tokens: Mapping[str, str] | None = None,
+    auth_subject: str = "",
+) -> str:
+    subject = auth_subject or authenticated_subject(headers, auth_token, auth_tokens or {})
+    if subject:
+        return f"auth:{subject}"
+    if trust_forwarded_for:
+        forwarded_for = header_value(headers, "X-Forwarded-For").split(",", 1)[0].strip()
+        forwarded_for_key = forwarded_for_client_key(forwarded_for)
+        if forwarded_for_key:
+            return forwarded_for_key
+    return remote_addr or "local"
+
+
+def request_id_from_headers(headers: Mapping[str, str]) -> str:
+    request_id = header_value(headers, "X-Request-ID")
+    if safe_request_id(request_id):
+        return request_id
+    return str(uuid4())
+
+
+def safe_request_id(value: str) -> bool:
+    if not value or len(value) > _MAX_REQUEST_ID_LENGTH:
+        return False
+    return all(33 <= ord(character) <= 126 for character in value)
+
+
+def safe_idempotency_key(value: str) -> bool:
+    if not value or len(value) > _MAX_IDEMPOTENCY_KEY_LENGTH:
+        return False
+    return all(33 <= ord(character) <= 126 for character in value)
+
+
+def safe_forwarded_for_client(value: str) -> bool:
+    return bool(forwarded_for_client_key(value))
+
+
+def forwarded_for_client_key(value: str) -> str:
+    if not safe_request_id(value):
+        return ""
+    try:
+        return str(ip_address(value))
+    except ValueError:
+        return ""
+
+
+def safe_trace_file_stem(value: Any) -> str:
+    raw_value = str(value or "trace")
+    safe_value = "".join(
+        character if character in _SAFE_TRACE_FILENAME_CHARS else "_"
+        for character in raw_value
+    ).strip("._")
+    return safe_value or "trace"
+
+
+def header_value(headers: Mapping[str, str], name: str) -> str:
+    expected = name.lower()
+    for key, value in headers.items():
+        if key.lower() == expected:
+            return value
+    return ""
+
+
+def safe_header_value(value: str) -> bool:
+    if not value:
+        return False
+    return all(32 <= ord(character) <= 126 for character in value)

package/src/kagent/service/server.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from typing import Optional, Type
+
+from kagent.service.active_runs import ActiveRunRegistry
+from kagent.service.runtime import (
+    ServiceConcurrencyLimiter,
+    ServiceConfig,
+    ServiceIdempotencyCache,
+    ServiceMetrics,
+    ServiceRateLimiter,
+    SqliteServiceIdempotencyCache,
+)
+from kagent.service.runtime_recovery import (
+    RuntimeInstanceLease,
+    reconcile_orphaned_runtime_traces,
+)
+
+
+class ProductionThreadingHTTPServer(ThreadingHTTPServer):
+    daemon_threads = False
+    block_on_close = True
+    allow_reuse_address = True
+    request_queue_size = 64
+
+    def server_close(self) -> None:
+        super().server_close()
+        lease = getattr(self, "service_runtime_instance_lease", None)
+        registry = getattr(self, "service_active_run_registry", None)
+        if isinstance(lease, RuntimeInstanceLease) and isinstance(
+            registry, ActiveRunRegistry
+        ):
+            lease.stop_when_idle(registry)
+
+
+def create_threading_server(
+    host: str,
+    port: int,
+    handler_class: Type[BaseHTTPRequestHandler],
+    *,
+    config: Optional[ServiceConfig] = None,
+) -> ThreadingHTTPServer:
+    server = ProductionThreadingHTTPServer((host, port), handler_class)
+    bound_host, bound_port = server.server_address
+    server.service_config = config or ServiceConfig(  # type: ignore[attr-defined]
+        host=str(bound_host),
+        port=int(bound_port),
+    )
+    service_metrics = ServiceMetrics()
+    server.service_metrics = service_metrics  # type: ignore[attr-defined]
+    server.service_rate_limiter = ServiceRateLimiter(  # type: ignore[attr-defined]
+        limit_per_minute=server.service_config.rate_limit_per_minute  # type: ignore[attr-defined]
+    )
+    server.service_concurrency_limiter = ServiceConcurrencyLimiter(  # type: ignore[attr-defined]
+        max_concurrent_runs=server.service_config.max_concurrent_runs  # type: ignore[attr-defined]
+    )
+    server.service_active_run_registry = ActiveRunRegistry()  # type: ignore[attr-defined]
+    if server.service_config.idempotency_cache_path:  # type: ignore[attr-defined]
+        server.service_idempotency_cache = SqliteServiceIdempotencyCache(  # type: ignore[attr-defined]
+            max_entries=server.service_config.idempotency_cache_size,  # type: ignore[attr-defined]
+            database_path=server.service_config.idempotency_cache_path,  # type: ignore[attr-defined]
+        )
+    else:
+        server.service_idempotency_cache = ServiceIdempotencyCache(  # type: ignore[attr-defined]
+            max_entries=server.service_config.idempotency_cache_size  # type: ignore[attr-defined]
+        )
+    server.service_runtime_instance_lease = None  # type: ignore[attr-defined]
+    server.service_runtime_reconciliation = {  # type: ignore[attr-defined]
+        "enabled": False,
+        "reason": "trace_persistence_disabled",
+    }
+    if server.service_config.trace_dir:  # type: ignore[attr-defined]
+        lease = RuntimeInstanceLease(
+            server.service_config.trace_dir,  # type: ignore[attr-defined]
+            instance_id=server.service_active_run_registry.instance_id,  # type: ignore[attr-defined]
+            heartbeat_seconds=(
+                server.service_config.runtime_instance_heartbeat_seconds  # type: ignore[attr-defined]
+            ),
+        )
+        server.service_runtime_instance_lease = lease  # type: ignore[attr-defined]
+        try:
+            lease.start()
+            server.service_runtime_reconciliation = (  # type: ignore[attr-defined]
+                reconcile_orphaned_runtime_traces(
+                    server.service_config.trace_dir,  # type: ignore[attr-defined]
+                    current_instance_id=(
+                        server.service_active_run_registry.instance_id  # type: ignore[attr-defined]
+                    ),
+                    stale_after_seconds=(
+                        server.service_config.runtime_orphaned_run_stale_seconds  # type: ignore[attr-defined]
+                    ),
+                )
+            )
+            service_metrics.record_runtime_reconciliation(
+                server.service_runtime_reconciliation  # type: ignore[attr-defined]
+            )
+        except OSError as exc:
+            lease.stop()
+            server.service_runtime_instance_lease = None  # type: ignore[attr-defined]
+            server.service_runtime_reconciliation = {  # type: ignore[attr-defined]
+                "enabled": False,
+                "reason": "trace_persistence_unavailable",
+                "error_type": type(exc).__name__,
+            }
+            service_metrics.record_runtime_reconciliation(
+                {"errors": 1},
+                status="unavailable",
+            )
+        except Exception:
+            lease.stop()
+            super(ProductionThreadingHTTPServer, server).server_close()
+            raise
+    return server

package/src/kagent/service/status.py

@@ -0,0 +1,233 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Callable, Dict, Optional
+from uuid import uuid4
+
+from kagent.integrations.backends import (
+    ExternalBackendConfig,
+    check_external_backends,
+)
+from kagent.providers.embeddings import EmbeddingProviderConfig
+from kagent.providers.llm import LLMProviderConfig
+from kagent.runtime.workspace import VIRTUAL_WORKSPACE_KINDS, RuntimeWorkspace
+from kagent.service import transport as service_transport
+from kagent.service.contract import service_openapi
+from kagent.service.errors import READINESS_FAILED
+from kagent.service.runtime import (
+    ServiceConfig,
+    SqliteServiceIdempotencyCache,
+)
+from kagent.service.trace_store import (
+    _ensure_owner_only_trace_dir,
+    _write_owner_only_temporary_trace,
+)
+
+
+def readiness_payload(config: Optional[ServiceConfig] = None) -> Dict[str, Any]:
+    active_config = config or ServiceConfig()
+    checks = {
+        "agent_config": _readiness_check(
+            _check_agent_config,
+            "agent_config_unavailable",
+        ),
+        "openapi": _readiness_check(service_openapi, "openapi_unavailable"),
+        "tools": _readiness_check(_check_tools, "tools_unavailable"),
+    }
+    if active_config.trace_dir:
+        checks["trace_persistence"] = _readiness_check(
+            lambda: _check_trace_persistence(active_config.trace_dir),
+            "trace_persistence_unavailable",
+        )
+    if active_config.runtime_workspace_dir:
+        checks["runtime_workspace"] = _readiness_check(
+            lambda: _check_runtime_workspace(active_config.runtime_workspace_dir),
+            "runtime_workspace_unavailable",
+        )
+    if active_config.idempotency_cache_path:
+        checks["idempotency_cache_persistence"] = _readiness_check(
+            lambda: _check_idempotency_cache_persistence(active_config),
+            "idempotency_cache_unavailable",
+        )
+    if _embedding_provider_configured(active_config):
+        checks["embedding_provider"] = _readiness_check(
+            lambda: _check_embedding_provider_config(active_config),
+            "embedding_provider_unavailable",
+        )
+    checks.update(check_external_backends(_external_backend_config(active_config)))
+    failed_checks = [name for name, value in checks.items() if value != "ok"]
+    status = "ready" if not failed_checks else "not_ready"
+    payload = {"status": status, "checks": checks, "failed_checks": failed_checks}
+    if failed_checks:
+        payload["error_code"] = READINESS_FAILED
+    return payload
+
+
+def service_config_snapshot(config: ServiceConfig) -> Dict[str, str]:
+    snapshot = {
+        "host": config.host,
+        "port": str(config.port),
+        "max_request_bytes": str(config.max_request_bytes),
+        "max_goal_chars": str(config.max_goal_chars),
+        "auth_required": str(config.auth_required).lower(),
+        "auth_subject_count": str(
+            len(config.auth_tokens) + (1 if config.auth_token else 0)
+        ),
+        "rate_limit_per_minute": str(config.rate_limit_per_minute),
+        "max_concurrent_runs": str(config.max_concurrent_runs),
+        "idempotency_cache_size": str(config.idempotency_cache_size),
+        "idempotency_cache_backend": (
+            "sqlite" if config.idempotency_cache_path else "memory"
+        ),
+        "idempotency_cache_path_configured": (
+            "true" if config.idempotency_cache_path else "false"
+        ),
+        "runtime_allowed_tools": _runtime_allowed_tools_snapshot(config),
+        "runtime_allowed_tools_by_subject_count": str(
+            len(config.runtime_allowed_tools_by_subject)
+        ),
+        "runtime_max_iterations": str(config.runtime_max_iterations),
+        "runtime_pending_approval_stale_seconds": str(
+            config.runtime_pending_approval_stale_seconds
+        ),
+        "runtime_instance_heartbeat_seconds": str(
+            config.runtime_instance_heartbeat_seconds
+        ),
+        "runtime_orphaned_run_stale_seconds": str(
+            config.runtime_orphaned_run_stale_seconds
+        ),
+        "allow_full_trace_response": str(config.allow_full_trace_response).lower(),
+        "protect_diagnostics": str(config.protect_diagnostics).lower(),
+        "trust_forwarded_for": str(config.trust_forwarded_for).lower(),
+        "run_timeout_seconds": str(config.run_timeout_seconds),
+        "request_timeout_seconds": str(config.request_timeout_seconds),
+        "trace_persistence": "enabled" if config.trace_dir else "disabled",
+        "runtime_workspace": (
+            "enabled" if config.runtime_workspace_dir else "disabled"
+        ),
+        "runtime_workspace_kinds": ",".join(VIRTUAL_WORKSPACE_KINDS),
+    }
+    snapshot.update(security_response_header_snapshot())
+    snapshot.update(trace_permission_policy_snapshot())
+    snapshot.update(_external_backend_config(config).redacted_snapshot())
+    snapshot.update(_embedding_provider_snapshot(config))
+    snapshot.update(llm_provider_snapshot())
+    return snapshot
+
+
+def trace_permission_policy_snapshot() -> Dict[str, str]:
+    return {
+        "trace_directory_permissions": "0700",
+        "trace_file_permissions": "0600",
+        "trace_probe_file_permissions": "0600",
+    }
+
+
+def llm_provider_snapshot() -> Dict[str, str]:
+    return LLMProviderConfig.from_sources().redacted_snapshot()
+
+
+def _embedding_provider_snapshot(config: ServiceConfig) -> Dict[str, str]:
+    return EmbeddingProviderConfig(
+        base_url=config.embedding_base_url,
+        api_key=config.embedding_api_key,
+        model=config.embedding_model,
+        timeout_seconds=config.embedding_timeout_seconds,
+        max_retries=config.embedding_max_retries,
+        retry_backoff_seconds=config.embedding_retry_backoff_seconds,
+    ).redacted_snapshot()
+
+
+def security_response_header_snapshot() -> Dict[str, str]:
+    return {
+        "security_response_headers": "enabled",
+        "cache_control_header": service_transport.CACHE_CONTROL_HEADER_VALUE,
+        "content_security_policy_header": (
+            service_transport.CONTENT_SECURITY_POLICY_HEADER_VALUE
+        ),
+        "referrer_policy_header": service_transport.REFERRER_POLICY_HEADER_VALUE,
+        "x_frame_options_header": service_transport.X_FRAME_OPTIONS_HEADER_VALUE,
+        "x_content_type_options_header": service_transport.NOSNIFF_HEADER_VALUE,
+    }
+
+
+def _runtime_allowed_tools_snapshot(config: ServiceConfig) -> str:
+    if not config.runtime_allowed_tools:
+        return "default"
+    return ",".join(config.runtime_allowed_tools)
+
+
+def _readiness_check(check: Callable[[], None], failure_code: str) -> str:
+    try:
+        check()
+    except Exception:  # pragma: no cover - defensive readiness detail
+        return f"failed: {failure_code}"
+    return "ok"
+
+
+def _check_agent_config() -> None:
+    from kagent.core.state import AgentConfig
+
+    AgentConfig()
+
+
+def _check_tools() -> None:
+    from kagent.core.tools import registered_tool_names
+
+    if not registered_tool_names():
+        raise RuntimeError("no tools registered")
+
+
+def _check_trace_persistence(trace_dir: str) -> None:
+    output_dir = Path(trace_dir)
+    _ensure_owner_only_trace_dir(output_dir)
+    probe_path = _write_owner_only_temporary_trace(
+        output_dir,
+        f"readiness-{uuid4().hex}.probe",
+        "ok\n",
+    )
+    probe_path.unlink()
+
+
+def _check_runtime_workspace(runtime_workspace_dir: str) -> None:
+    RuntimeWorkspace(runtime_workspace_dir).ensure_layout()
+
+
+def _check_idempotency_cache_persistence(config: ServiceConfig) -> None:
+    SqliteServiceIdempotencyCache(
+        max_entries=config.idempotency_cache_size,
+        database_path=config.idempotency_cache_path,
+    ).snapshot()
+
+
+def _embedding_provider_configured(config: ServiceConfig) -> bool:
+    return bool(
+        config.embedding_base_url
+        or config.embedding_api_key
+        or config.embedding_model
+    )
+
+
+def _check_embedding_provider_config(config: ServiceConfig) -> None:
+    if not config.embedding_base_url:
+        raise ValueError("embedding base_url is required")
+    if not config.embedding_model:
+        raise ValueError("embedding model is required")
+    EmbeddingProviderConfig(
+        base_url=config.embedding_base_url,
+        api_key=config.embedding_api_key,
+        model=config.embedding_model,
+        timeout_seconds=config.embedding_timeout_seconds,
+        max_retries=config.embedding_max_retries,
+        retry_backoff_seconds=config.embedding_retry_backoff_seconds,
+    )
+
+
+def _external_backend_config(config: ServiceConfig) -> ExternalBackendConfig:
+    return ExternalBackendConfig(
+        redis_url=config.redis_url,
+        milvus_url=config.milvus_url,
+        kafka_audit_url=config.kafka_audit_url,
+        kafka_audit_topic=config.kafka_audit_topic,
+        timeout_seconds=config.external_backend_timeout_seconds,
+    )