mirror of
https://github.com/langgenius/dify.git
synced 2026-07-20 09:38:32 +08:00
feat: daily sync (#38593)
Co-authored-by: 盐粒 Yanli <mail@yanli.one> Co-authored-by: yyh <yuanyouhuilyz@gmail.com> Co-authored-by: Joel <iamjoel007@gmail.com> Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: QuantumGhost <obelisk.reg+git@gmail.com> Co-authored-by: yyh <92089059+lyzno1@users.noreply.github.com> Co-authored-by: 盐粒 Yanli <yanli@dify.ai>
This commit is contained in:
parent
9252d81826
commit
925f97be20
@ -663,6 +663,9 @@ PLUGIN_MODEL_SCHEMA_CACHE_TTL=3600
|
||||
PLUGIN_MODEL_PROVIDERS_CACHE_TTL=86400
|
||||
INNER_API_KEY_FOR_PLUGIN=QaHbTe77CtuXmsfyhR7+vRjI/+XbV1AaFy691iy+kGDv2Jvy0/eAh8Y1
|
||||
|
||||
# Dify Agent backend
|
||||
AGENT_BACKEND_BASE_URL=http://localhost:5050
|
||||
|
||||
# Marketplace configuration
|
||||
MARKETPLACE_ENABLED=true
|
||||
MARKETPLACE_API_URL=https://marketplace.dify.ai
|
||||
|
||||
@ -25,11 +25,10 @@ class AgentBackendConfig(BaseSettings):
|
||||
AGENT_SHELL_ENABLED: bool = Field(
|
||||
description=(
|
||||
"Inject the dify.shell layer (sandboxed bash workspace) into Agent runs. "
|
||||
"Requires the agent backend to be wired with a shellctl entrypoint; keep it "
|
||||
"off until shellctl is deployed, otherwise every agent run that includes the "
|
||||
"shell layer will fail."
|
||||
"Requires the agent backend to be wired with a shellctl entrypoint before "
|
||||
"shell-using Agent runs are executed."
|
||||
),
|
||||
default=False,
|
||||
default=True,
|
||||
)
|
||||
|
||||
AGENT_APP_TEXT_DELTA_DEBOUNCE_SECONDS: NonNegativeFloat = Field(
|
||||
|
||||
@ -363,7 +363,10 @@ class FileAccessConfig(BaseSettings):
|
||||
INTERNAL_FILES_URL: str = Field(
|
||||
description="Internal base URL for file access within Docker network,"
|
||||
" used for plugin daemon and internal service communication."
|
||||
" Falls back to FILES_URL if not specified.",
|
||||
" Explicit INTERNAL_FILES_URL takes precedence; otherwise SERVER_CONSOLE_API_URL is used,"
|
||||
" then FILES_URL.",
|
||||
validation_alias=AliasChoices("INTERNAL_FILES_URL", "SERVER_CONSOLE_API_URL"),
|
||||
alias_priority=1,
|
||||
default="",
|
||||
)
|
||||
|
||||
|
||||
@ -1,5 +1,3 @@
|
||||
from mimetypes import guess_extension
|
||||
|
||||
from flask import request
|
||||
from flask_restx import Resource
|
||||
from flask_restx.api import HTTPStatus
|
||||
@ -8,7 +6,7 @@ from werkzeug.exceptions import Forbidden
|
||||
|
||||
import services
|
||||
from core.tools.signature import verify_plugin_file_signature
|
||||
from core.tools.tool_file_manager import ToolFileManager
|
||||
from core.tools.tool_file_manager import ToolFileManager, resolve_extension
|
||||
from core.workflow.file_reference import build_file_reference
|
||||
from fields.file_fields import FileResponse
|
||||
|
||||
@ -110,7 +108,7 @@ class PluginUploadFileApi(Resource):
|
||||
conversation_id=args.conversation_id,
|
||||
)
|
||||
|
||||
extension = guess_extension(tool_file.mimetype) or ".bin"
|
||||
extension = resolve_extension(filename=tool_file.name, mimetype=tool_file.mimetype)
|
||||
preview_url = ToolFileManager.sign_file(tool_file_id=tool_file.id, extension=extension)
|
||||
|
||||
# Create a dictionary with all the necessary attributes
|
||||
|
||||
@ -476,6 +476,7 @@ class PluginDownloadFileRequestApi(Resource):
|
||||
user_from=payload.user_from,
|
||||
invoke_from=payload.invoke_from,
|
||||
file_mapping=payload.file.model_dump(mode="python", exclude_none=True),
|
||||
for_external=payload.for_external,
|
||||
)
|
||||
return BaseBackwardsInvocationResponse(
|
||||
data={
|
||||
|
||||
@ -27,6 +27,7 @@ from clients.agent_backend import (
|
||||
AgentBackendInternalEventType,
|
||||
AgentBackendRunClient,
|
||||
AgentBackendRunEventAdapter,
|
||||
AgentBackendRunFailedInternalEvent,
|
||||
AgentBackendRunSucceededInternalEvent,
|
||||
AgentBackendStreamInternalEvent,
|
||||
extract_runtime_layer_specs,
|
||||
@ -57,6 +58,14 @@ from core.workflow.nodes.agent_v2.ask_human_resume import build_deferred_tool_re
|
||||
from extensions.ext_database import db
|
||||
from graphon.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage
|
||||
from graphon.model_runtime.entities.message_entities import AssistantPromptMessage, PromptMessage, UserPromptMessage
|
||||
from graphon.model_runtime.errors.invoke import (
|
||||
InvokeAuthorizationError,
|
||||
InvokeBadRequestError,
|
||||
InvokeConnectionError,
|
||||
InvokeError,
|
||||
InvokeRateLimitError,
|
||||
InvokeServerUnavailableError,
|
||||
)
|
||||
from models.agent_config_entities import AgentSoulConfig
|
||||
from models.enums import CreatorUserRole
|
||||
from models.model import MessageAgentThought
|
||||
@ -71,6 +80,22 @@ class _DefaultSessionScopeSnapshotId:
|
||||
|
||||
_DEFAULT_SESSION_SCOPE_SNAPSHOT_ID = _DefaultSessionScopeSnapshotId()
|
||||
|
||||
_AGENT_BACKEND_INVOKE_ERROR_BY_REASON: Mapping[str, type[InvokeError]] = {
|
||||
"InvokeAuthorizationError": InvokeAuthorizationError,
|
||||
"InvokeBadRequestError": InvokeBadRequestError,
|
||||
"CredentialsValidateFailedError": InvokeBadRequestError,
|
||||
"InvokeConnectionError": InvokeConnectionError,
|
||||
"InvokeRateLimitError": InvokeRateLimitError,
|
||||
"InvokeServerUnavailableError": InvokeServerUnavailableError,
|
||||
}
|
||||
|
||||
|
||||
def _agent_backend_failure_to_exception(event: AgentBackendRunFailedInternalEvent) -> Exception:
|
||||
err_cls = _AGENT_BACKEND_INVOKE_ERROR_BY_REASON.get(event.reason or "")
|
||||
if err_cls is not None:
|
||||
return err_cls(event.error)
|
||||
return AgentBackendError(event.error or "Agent backend run did not complete successfully.")
|
||||
|
||||
|
||||
def _prompt_messages_from_query(user_query: str | None) -> list[PromptMessage]:
|
||||
if not user_query:
|
||||
@ -412,12 +437,15 @@ class _AgentProcessRecorder:
|
||||
def _lookup_tool_thought(self, *, index: int, tool_call_id: str | None) -> str | None:
|
||||
if tool_call_id and tool_call_id in self._tool_by_call_id:
|
||||
return self._tool_by_call_id[tool_call_id]
|
||||
if index < 0:
|
||||
return None
|
||||
return self._tool_by_index.get(index)
|
||||
|
||||
def _remember_tool_thought(
|
||||
self, *, index: int, tool_call_id: str | None, tool_name: str | None, thought_id: str
|
||||
) -> None:
|
||||
self._tool_by_index[index] = thought_id
|
||||
if index >= 0:
|
||||
self._tool_by_index[index] = thought_id
|
||||
if tool_call_id:
|
||||
self._tool_by_call_id[tool_call_id] = thought_id
|
||||
if tool_name:
|
||||
@ -433,6 +461,10 @@ class _AgentProcessRecorder:
|
||||
return None
|
||||
|
||||
def _mark_tool_observed(self, thought_id: str) -> None:
|
||||
self._tool_by_index = {index: value for index, value in self._tool_by_index.items() if value != thought_id}
|
||||
self._tool_by_call_id = {
|
||||
tool_call_id: value for tool_call_id, value in self._tool_by_call_id.items() if value != thought_id
|
||||
}
|
||||
for open_thought_ids in self._open_tool_by_name.values():
|
||||
open_thought_ids.discard(thought_id)
|
||||
|
||||
@ -530,7 +562,12 @@ def _event_index(data: dict[str, Any]) -> int:
|
||||
|
||||
|
||||
def _string_or_none(value: Any) -> str | None:
|
||||
return value if isinstance(value, str) and value else None
|
||||
if not isinstance(value, str):
|
||||
return None
|
||||
normalized = value.strip()
|
||||
if not normalized or normalized.lower() in {"none", "null"}:
|
||||
return None
|
||||
return normalized
|
||||
|
||||
|
||||
def _json_or_text(value: Any) -> str | None:
|
||||
@ -652,8 +689,9 @@ class AgentAppRunner:
|
||||
return
|
||||
|
||||
if not isinstance(terminal, AgentBackendRunSucceededInternalEvent):
|
||||
error = getattr(terminal, "error", None) or "Agent backend run did not complete successfully."
|
||||
raise AgentBackendError(str(error))
|
||||
if isinstance(terminal, AgentBackendRunFailedInternalEvent):
|
||||
raise _agent_backend_failure_to_exception(terminal)
|
||||
raise AgentBackendError("Agent backend run did not complete successfully.")
|
||||
|
||||
answer = self._terminal_output_to_answer(terminal.output)
|
||||
try:
|
||||
|
||||
@ -8,7 +8,7 @@ from pydantic import JsonValue
|
||||
from core.app.entities.app_invoke_entities import InvokeFrom
|
||||
from core.app.entities.task_entities import AppBlockingResponse, AppStreamResponse
|
||||
from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError
|
||||
from graphon.model_runtime.errors.invoke import InvokeError
|
||||
from graphon.model_runtime.errors.invoke import InvokeError, InvokeRateLimitError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@ -127,6 +127,7 @@ class AppGenerateResponseConverter[TBlockingResponse: AppBlockingResponse](ABC):
|
||||
},
|
||||
ModelCurrentlyNotSupportError: {"code": "model_currently_not_support", "status": 400},
|
||||
InvokeError: {"code": "completion_request_error", "status": 400},
|
||||
InvokeRateLimitError: {"code": "rate_limit_error", "status": 429},
|
||||
}
|
||||
|
||||
# Determine the response based on the type of exception
|
||||
|
||||
@ -420,6 +420,7 @@ class EasyUIBasedGenerateTaskPipeline(BasedGenerateTaskPipeline[EasyUIAppGenerat
|
||||
message.total_price = usage.total_price
|
||||
message.currency = usage.currency
|
||||
self._task_state.llm_result.usage.latency = message.provider_response_latency
|
||||
self._task_state.metadata.usage = self._task_state.llm_result.usage
|
||||
message.message_metadata = self._task_state.metadata.model_dump_json()
|
||||
|
||||
if trace_manager:
|
||||
|
||||
@ -276,6 +276,7 @@ class RequestRequestDownloadFile(BaseModel):
|
||||
"validation",
|
||||
]
|
||||
file: RequestDownloadFileMapping
|
||||
for_external: bool = True
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
|
||||
@ -4,6 +4,7 @@ import hmac
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import urllib.parse
|
||||
from collections.abc import Generator
|
||||
from mimetypes import guess_extension, guess_type
|
||||
from uuid import uuid4
|
||||
@ -26,7 +27,7 @@ logger = logging.getLogger(__name__)
|
||||
class ToolFileManager:
|
||||
@staticmethod
|
||||
def _build_graph_file_reference(tool_file: ToolFile) -> File:
|
||||
extension = guess_extension(tool_file.mimetype) or ".bin"
|
||||
extension = resolve_extension(filename=tool_file.name, mimetype=tool_file.mimetype)
|
||||
return File(
|
||||
file_type=get_file_type_by_mime_type(tool_file.mimetype),
|
||||
transfer_method=FileTransferMethod.TOOL_FILE,
|
||||
@ -70,7 +71,7 @@ class ToolFileManager:
|
||||
mimetype: str,
|
||||
filename: str | None = None,
|
||||
) -> ToolFile:
|
||||
extension = guess_extension(mimetype) or ".bin"
|
||||
extension = resolve_extension(filename=filename, mimetype=mimetype)
|
||||
unique_name = uuid4().hex
|
||||
unique_filename = f"{unique_name}{extension}"
|
||||
# default just as before
|
||||
@ -120,7 +121,8 @@ class ToolFileManager:
|
||||
or response.headers.get("Content-Type", "").split(";")[0].strip()
|
||||
or "application/octet-stream"
|
||||
)
|
||||
extension = guess_extension(mimetype) or ".bin"
|
||||
url_filename = os.path.basename(urllib.parse.urlparse(file_url).path)
|
||||
extension = resolve_extension(filename=url_filename, mimetype=mimetype)
|
||||
unique_name = uuid4().hex
|
||||
filename = f"{unique_name}{extension}"
|
||||
filepath = f"tools/{tenant_id}/{filename}"
|
||||
@ -220,4 +222,11 @@ def _factory() -> ToolFileManager:
|
||||
return ToolFileManager()
|
||||
|
||||
|
||||
def resolve_extension(*, filename: str | None, mimetype: str) -> str:
|
||||
filename_extension = os.path.splitext(filename or "")[1].lower()
|
||||
if filename_extension:
|
||||
return filename_extension
|
||||
return guess_extension(mimetype) or ".bin"
|
||||
|
||||
|
||||
set_tool_file_manager_factory(_factory)
|
||||
|
||||
@ -3,7 +3,6 @@ import re
|
||||
from collections.abc import Generator
|
||||
from datetime import date, datetime
|
||||
from decimal import Decimal
|
||||
from mimetypes import guess_extension
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
@ -11,7 +10,7 @@ import numpy as np
|
||||
import pytz
|
||||
|
||||
from core.tools.entities.tool_entities import ToolInvokeMessage
|
||||
from core.tools.tool_file_manager import ToolFileManager
|
||||
from core.tools.tool_file_manager import ToolFileManager, resolve_extension
|
||||
from core.workflow.file_reference import parse_file_reference
|
||||
from graphon.file import File, FileTransferMethod, FileType
|
||||
from libs.login import current_user
|
||||
@ -91,7 +90,8 @@ class ToolFileMessageTransformer:
|
||||
conversation_id=conversation_id,
|
||||
)
|
||||
|
||||
url = f"/files/tools/{tool_file.id}{guess_extension(tool_file.mimetype) or '.png'}"
|
||||
extension = resolve_extension(filename=tool_file.name, mimetype=tool_file.mimetype)
|
||||
url = cls.get_tool_file_url(tool_file_id=tool_file.id, extension=extension)
|
||||
meta = cls._with_tool_file_meta(
|
||||
message.meta,
|
||||
tool_file_id=str(tool_file.id),
|
||||
@ -136,7 +136,8 @@ class ToolFileMessageTransformer:
|
||||
filename=filename,
|
||||
)
|
||||
|
||||
url = cls.get_tool_file_url(tool_file_id=tool_file.id, extension=guess_extension(tool_file.mimetype))
|
||||
extension = resolve_extension(filename=tool_file.name, mimetype=tool_file.mimetype)
|
||||
url = cls.get_tool_file_url(tool_file_id=tool_file.id, extension=extension)
|
||||
meta = cls._with_tool_file_meta(meta, tool_file_id=str(tool_file.id))
|
||||
|
||||
# check if file is image
|
||||
|
||||
@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, override
|
||||
from agenton.compositor import CompositorSessionSnapshot
|
||||
|
||||
from clients.agent_backend import (
|
||||
AgentBackendAgentMessageDeltaInternalEvent,
|
||||
AgentBackendDeferredToolCallInternalEvent,
|
||||
AgentBackendError,
|
||||
AgentBackendHTTPError,
|
||||
@ -481,6 +482,10 @@ class DifyAgentNode(Node[DifyAgentNodeData]):
|
||||
if isinstance(internal_event, AgentBackendStreamInternalEvent):
|
||||
self._record_stream_metadata(metadata, internal_event)
|
||||
continue
|
||||
if internal_event.type == AgentBackendInternalEventType.AGENT_MESSAGE_DELTA:
|
||||
if isinstance(internal_event, AgentBackendAgentMessageDeltaInternalEvent):
|
||||
self._record_agent_message_delta_metadata(metadata, internal_event)
|
||||
continue
|
||||
metadata["agent_backend"] = {
|
||||
**dict(metadata.get("agent_backend") or {}),
|
||||
"stream_event_count": stream_event_count,
|
||||
@ -734,6 +739,17 @@ class DifyAgentNode(Node[DifyAgentNodeData]):
|
||||
agent_backend["usage"] = dict(usage)
|
||||
metadata["agent_backend"] = agent_backend
|
||||
|
||||
@staticmethod
|
||||
def _record_agent_message_delta_metadata(
|
||||
metadata: dict[str, Any], event: AgentBackendAgentMessageDeltaInternalEvent
|
||||
) -> None:
|
||||
agent_backend = dict(metadata.get("agent_backend") or {})
|
||||
agent_backend["agent_message_delta_count"] = int(agent_backend.get("agent_message_delta_count") or 0) + 1
|
||||
agent_backend["agent_message_delta_length"] = int(agent_backend.get("agent_message_delta_length") or 0) + len(
|
||||
event.delta
|
||||
)
|
||||
metadata["agent_backend"] = agent_backend
|
||||
|
||||
@classmethod
|
||||
@override
|
||||
def _extract_variable_selector_to_variable_mapping(
|
||||
|
||||
@ -10,13 +10,13 @@ trustworthy metadata.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from mimetypes import guess_extension
|
||||
from uuid import UUID
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.exc import DataError, SQLAlchemyError
|
||||
|
||||
from core.db.session_factory import session_factory
|
||||
from core.tools.tool_file_manager import resolve_extension
|
||||
from core.workflow.file_reference import build_file_reference
|
||||
from graphon.file import File, FileTransferMethod, get_file_type_by_mime_type
|
||||
from models.tools import ToolFile
|
||||
@ -46,7 +46,7 @@ def reback_tool_file_output(*, tenant_id: str, tool_file_id: str) -> File | None
|
||||
return None
|
||||
|
||||
mime_type = tool_file.mimetype or ""
|
||||
extension = guess_extension(mime_type) or ".bin"
|
||||
extension = resolve_extension(filename=tool_file.name, mimetype=mime_type)
|
||||
return File(
|
||||
type=get_file_type_by_mime_type(mime_type),
|
||||
transfer_method=FileTransferMethod.TOOL_FILE,
|
||||
|
||||
@ -3,6 +3,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import mimetypes
|
||||
import os
|
||||
import uuid
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any, Literal, NotRequired, TypedDict, assert_never, cast
|
||||
@ -285,7 +286,7 @@ def _build_from_remote_url(
|
||||
raise ValueError("Invalid file url")
|
||||
|
||||
mime_type, filename, file_size = get_remote_file_info(url)
|
||||
extension = mimetypes.guess_extension(mime_type) or ("." + filename.split(".")[-1] if "." in filename else ".bin")
|
||||
extension = os.path.splitext(filename)[1].lower() or mimetypes.guess_extension(mime_type) or ".bin"
|
||||
detected_file_type = standardize_file_type(extension=extension, mime_type=mime_type)
|
||||
file_type = _resolve_file_type(
|
||||
detected_file_type=detected_file_type,
|
||||
@ -326,7 +327,12 @@ def _build_from_tool_file(
|
||||
if tool_file is None:
|
||||
raise ValueError(f"ToolFile {tool_file_id} not found")
|
||||
|
||||
extension = "." + tool_file.file_key.split(".")[-1] if "." in tool_file.file_key else ".bin"
|
||||
extension = (
|
||||
os.path.splitext(tool_file.name)[1].lower()
|
||||
or mimetypes.guess_extension(tool_file.mimetype)
|
||||
or os.path.splitext(tool_file.file_key)[1].lower()
|
||||
or ".bin"
|
||||
)
|
||||
detected_file_type = standardize_file_type(extension=extension, mime_type=tool_file.mimetype)
|
||||
file_type = _resolve_file_type(
|
||||
detected_file_type=detected_file_type,
|
||||
|
||||
@ -1,9 +1,10 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
from decimal import Decimal
|
||||
from uuid import uuid4
|
||||
|
||||
from pydantic import Field, field_validator
|
||||
from pydantic import Field, computed_field, field_validator
|
||||
|
||||
from core.entities.execution_extra_content import ExecutionExtraContentDomainModel
|
||||
from fields.base import ResponseModel
|
||||
@ -55,10 +56,19 @@ class MessageListItem(ResponseModel):
|
||||
created_at: int | None = None
|
||||
agent_thoughts: list[AgentThought]
|
||||
message_files: list[MessageFile]
|
||||
message_tokens: int = 0
|
||||
answer_tokens: int = 0
|
||||
provider_response_latency: float = 0
|
||||
total_price: Decimal | None = None
|
||||
currency: str | None = None
|
||||
status: str
|
||||
error: str | None = None
|
||||
extra_contents: list[ExecutionExtraContentDomainModel]
|
||||
|
||||
@computed_field
|
||||
def total_tokens(self) -> int:
|
||||
return self.message_tokens + self.answer_tokens
|
||||
|
||||
@field_validator("inputs", mode="before")
|
||||
@classmethod
|
||||
def _normalize_inputs(cls, value: JSONValueType) -> JSONValueType:
|
||||
|
||||
@ -17614,19 +17614,25 @@ Built-in tool icons are URL strings; API-based tool icons are provider-defined p
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| agent_thoughts | [ [AgentThought](#agentthought) ] | | Yes |
|
||||
| answer | string | | Yes |
|
||||
| answer_tokens | integer | | No |
|
||||
| conversation_id | string | | Yes |
|
||||
| created_at | integer | | No |
|
||||
| currency | string | | No |
|
||||
| error | string | | No |
|
||||
| extra_contents | [ [HumanInputContent](#humaninputcontent) ] | | Yes |
|
||||
| feedback | [SimpleFeedback](#simplefeedback) | | No |
|
||||
| id | string | | Yes |
|
||||
| inputs | object | | Yes |
|
||||
| message_files | [ [MessageFile](#messagefile) ] | | Yes |
|
||||
| message_tokens | integer | | No |
|
||||
| metadata | [JSONValueType](#jsonvaluetype) | | No |
|
||||
| parent_message_id | string | | No |
|
||||
| provider_response_latency | number | | No |
|
||||
| query | string | | Yes |
|
||||
| retriever_resources | [ [RetrieverResource](#retrieverresource) ] | | Yes |
|
||||
| status | string | | Yes |
|
||||
| total_price | string | | No |
|
||||
| total_tokens | integer | | Yes |
|
||||
|
||||
#### ExternalApiTemplateListQuery
|
||||
|
||||
|
||||
@ -3467,18 +3467,24 @@ Model class for i18n object.
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| agent_thoughts | [ [AgentThought](#agentthought) ] | | Yes |
|
||||
| answer | string | | Yes |
|
||||
| answer_tokens | integer | | No |
|
||||
| conversation_id | string | | Yes |
|
||||
| created_at | integer | | No |
|
||||
| currency | string | | No |
|
||||
| error | string | | No |
|
||||
| extra_contents | [ [HumanInputContent](#humaninputcontent) ] | | Yes |
|
||||
| feedback | [SimpleFeedback](#simplefeedback) | | No |
|
||||
| id | string | | Yes |
|
||||
| inputs | object | | Yes |
|
||||
| message_files | [ [MessageFile](#messagefile) ] | | Yes |
|
||||
| message_tokens | integer | | No |
|
||||
| parent_message_id | string | | No |
|
||||
| provider_response_latency | number | | No |
|
||||
| query | string | | Yes |
|
||||
| retriever_resources | [ [RetrieverResource](#retrieverresource) ] | | Yes |
|
||||
| status | string | | Yes |
|
||||
| total_price | string | | No |
|
||||
| total_tokens | integer | | Yes |
|
||||
|
||||
#### MessageListQuery
|
||||
|
||||
|
||||
@ -1685,19 +1685,25 @@ in form definiton, or a variable while the workflow is running.
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| agent_thoughts | [ [AgentThought](#agentthought) ] | | Yes |
|
||||
| answer | string | | Yes |
|
||||
| answer_tokens | integer | | No |
|
||||
| conversation_id | string | | Yes |
|
||||
| created_at | integer | | No |
|
||||
| currency | string | | No |
|
||||
| error | string | | No |
|
||||
| extra_contents | [ [HumanInputContent](#humaninputcontent) ] | | Yes |
|
||||
| feedback | [SimpleFeedback](#simplefeedback) | | No |
|
||||
| id | string | | Yes |
|
||||
| inputs | object | | Yes |
|
||||
| message_files | [ [MessageFile](#messagefile) ] | | Yes |
|
||||
| message_tokens | integer | | No |
|
||||
| metadata | [JSONValueType](#jsonvaluetype) | | No |
|
||||
| parent_message_id | string | | No |
|
||||
| provider_response_latency | number | | No |
|
||||
| query | string | | Yes |
|
||||
| retriever_resources | [ [RetrieverResource](#retrieverresource) ] | | Yes |
|
||||
| status | string | | Yes |
|
||||
| total_price | string | | No |
|
||||
| total_tokens | integer | | Yes |
|
||||
|
||||
#### WebModelConfigResponse
|
||||
|
||||
|
||||
@ -45,6 +45,7 @@ class FileRequestService:
|
||||
user_from: UserFrom | str,
|
||||
invoke_from: InvokeFrom | str,
|
||||
file_mapping: Mapping[str, Any],
|
||||
for_external: bool = True,
|
||||
) -> DownloadFileRequestResult:
|
||||
"""Resolve one file mapping into signed download metadata.
|
||||
|
||||
@ -61,7 +62,7 @@ class FileRequestService:
|
||||
)
|
||||
with bind_file_access_scope(scope):
|
||||
file = self._build_file(mapping=file_mapping, tenant_id=tenant_id)
|
||||
download_url = file_helpers.resolve_file_url(file, for_external=True)
|
||||
download_url = file_helpers.resolve_file_url(file, for_external=for_external)
|
||||
|
||||
if not download_url:
|
||||
raise ValueError("file does not support signed download")
|
||||
|
||||
@ -75,6 +75,7 @@ def test_dify_config(monkeypatch: pytest.MonkeyPatch):
|
||||
# default values
|
||||
assert config.EDITION == "SELF_HOSTED"
|
||||
assert config.API_COMPRESSION_ENABLED is False
|
||||
assert config.AGENT_SHELL_ENABLED is True
|
||||
assert config.SENTRY_TRACES_SAMPLE_RATE == 1.0
|
||||
assert config.TEMPLATE_TRANSFORM_MAX_LENGTH == 400_000
|
||||
|
||||
@ -110,6 +111,25 @@ def test_http_timeout_defaults(monkeypatch: pytest.MonkeyPatch):
|
||||
assert config.HTTP_REQUEST_MAX_WRITE_TIMEOUT == 600
|
||||
|
||||
|
||||
def test_internal_files_url_falls_back_to_server_console_api_url(monkeypatch: pytest.MonkeyPatch):
|
||||
os.environ.clear()
|
||||
monkeypatch.setenv("SERVER_CONSOLE_API_URL", "http://api:5001")
|
||||
|
||||
config = DifyConfig(_env_file=None)
|
||||
|
||||
assert config.INTERNAL_FILES_URL == "http://api:5001"
|
||||
|
||||
|
||||
def test_internal_files_url_prefers_explicit_value(monkeypatch: pytest.MonkeyPatch):
|
||||
os.environ.clear()
|
||||
monkeypatch.setenv("INTERNAL_FILES_URL", "http://files-internal:5001")
|
||||
monkeypatch.setenv("SERVER_CONSOLE_API_URL", "http://api:5001")
|
||||
|
||||
config = DifyConfig(_env_file=None)
|
||||
|
||||
assert config.INTERNAL_FILES_URL == "http://files-internal:5001"
|
||||
|
||||
|
||||
# NOTE: If there is a `.env` file in your Workspace, this test might not succeed as expected.
|
||||
# This is due to `pymilvus` loading all the variables from the `.env` file into `os.environ`.
|
||||
def test_flask_configs(monkeypatch: pytest.MonkeyPatch):
|
||||
|
||||
@ -49,6 +49,8 @@ def make_message():
|
||||
msg.query = "hello"
|
||||
msg.re_sign_file_url_answer = ""
|
||||
msg.user_feedback = MagicMock(rating=None)
|
||||
msg.total_price = None
|
||||
msg.currency = None
|
||||
msg.status = "normal"
|
||||
msg.error = None
|
||||
return msg
|
||||
|
||||
@ -34,11 +34,11 @@ class DummyFile:
|
||||
|
||||
|
||||
class DummyToolFile:
|
||||
def __init__(self):
|
||||
def __init__(self, name="test.txt", mimetype="text/plain"):
|
||||
self.id = "file-id"
|
||||
self.name = "test.txt"
|
||||
self.name = name
|
||||
self.size = 10
|
||||
self.mimetype = "text/plain"
|
||||
self.mimetype = mimetype
|
||||
self.original_url = "http://original"
|
||||
self.user_id = "user-1"
|
||||
self.tenant_id = "tenant-1"
|
||||
@ -56,7 +56,7 @@ class TestPluginUploadFileApi:
|
||||
mock_get_user,
|
||||
mock_verify_signature,
|
||||
):
|
||||
dummy_file = DummyFile()
|
||||
dummy_file = DummyFile(filename="report.docx", mimetype="application/octet-stream")
|
||||
|
||||
module.request = fake_request(
|
||||
{
|
||||
@ -71,7 +71,10 @@ class TestPluginUploadFileApi:
|
||||
)
|
||||
|
||||
tool_file_manager_instance = mock_tool_file_manager.return_value
|
||||
tool_file_manager_instance.create_file_by_raw.return_value = DummyToolFile()
|
||||
tool_file_manager_instance.create_file_by_raw.return_value = DummyToolFile(
|
||||
name="report.docx",
|
||||
mimetype="application/octet-stream",
|
||||
)
|
||||
|
||||
mock_tool_file_manager.sign_file.return_value = "signed-url"
|
||||
|
||||
@ -84,10 +87,12 @@ class TestPluginUploadFileApi:
|
||||
assert result["id"] == "file-id"
|
||||
assert result["reference"] == build_file_reference(record_id="file-id")
|
||||
assert result["preview_url"] == "signed-url"
|
||||
assert result["extension"] == ".docx"
|
||||
mock_verify_signature.assert_called_once()
|
||||
assert mock_verify_signature.call_args.kwargs["conversation_id"] == "conversation-1"
|
||||
tool_file_manager_instance.create_file_by_raw.assert_called_once()
|
||||
assert tool_file_manager_instance.create_file_by_raw.call_args.kwargs["conversation_id"] == "conversation-1"
|
||||
mock_tool_file_manager.sign_file.assert_called_once_with(tool_file_id="file-id", extension=".docx")
|
||||
|
||||
def test_missing_file(self):
|
||||
module.request = fake_request(
|
||||
|
||||
@ -318,6 +318,7 @@ class TestPluginDownloadFileRequestApi:
|
||||
mock_payload.user_id = "user-id"
|
||||
mock_payload.user_from = "account"
|
||||
mock_payload.invoke_from = "debugger"
|
||||
mock_payload.for_external = False
|
||||
reference = build_file_reference(record_id="tool-file-1")
|
||||
mock_payload.file.model_dump.return_value = {
|
||||
"transfer_method": "tool_file",
|
||||
@ -333,6 +334,7 @@ class TestPluginDownloadFileRequestApi:
|
||||
user_from="account",
|
||||
invoke_from="debugger",
|
||||
file_mapping={"transfer_method": "tool_file", "reference": reference},
|
||||
for_external=False,
|
||||
)
|
||||
assert result["data"] == {
|
||||
"filename": "report.pdf",
|
||||
|
||||
@ -37,6 +37,7 @@ from pydantic_ai.messages import (
|
||||
from clients.agent_backend import (
|
||||
AgentBackendError,
|
||||
AgentBackendRunEventAdapter,
|
||||
AgentBackendRunFailedInternalEvent,
|
||||
AgentBackendStreamInternalEvent,
|
||||
FakeAgentBackendRunClient,
|
||||
FakeAgentBackendScenario,
|
||||
@ -54,6 +55,7 @@ from core.app.entities.queue_entities import (
|
||||
QueueMessageEndEvent,
|
||||
)
|
||||
from core.workflow.nodes.agent_v2.ask_human_resume import AskHumanResumeOutcome
|
||||
from graphon.model_runtime.errors.invoke import InvokeRateLimitError
|
||||
from models.agent_config_entities import AgentSoulConfig
|
||||
from models.model import MessageAgentThought
|
||||
|
||||
@ -1039,6 +1041,130 @@ def test_tool_result_without_call_id_matches_unique_open_tool_name(monkeypatch):
|
||||
assert rows[0].observation == "Knowledge base search results: browser skill"
|
||||
|
||||
|
||||
def test_repeated_tool_calls_without_call_id_or_index_create_distinct_rows(monkeypatch):
|
||||
fake_session = _FakeDbSession()
|
||||
monkeypatch.setattr(app_runner_module.db, "session", fake_session)
|
||||
qm = _FakeQueueManager()
|
||||
recorder = app_runner_module._AgentProcessRecorder(
|
||||
dify_context=_dify_ctx(),
|
||||
message_id="msg-1",
|
||||
queue_manager=qm, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
recorder.handle_stream_event(
|
||||
AgentBackendStreamInternalEvent(
|
||||
run_id="run-1",
|
||||
data={
|
||||
"event_kind": "function_tool_call",
|
||||
"part": {
|
||||
"part_kind": "tool-call",
|
||||
"tool_name": "shell_run",
|
||||
"args": {"script": "lookup find"},
|
||||
},
|
||||
},
|
||||
)
|
||||
)
|
||||
recorder.handle_stream_event(
|
||||
AgentBackendStreamInternalEvent(
|
||||
run_id="run-1",
|
||||
data={
|
||||
"event_kind": "function_tool_result",
|
||||
"part": {
|
||||
"part_kind": "tool-return",
|
||||
"tool_name": "shell_run",
|
||||
"content": "find output",
|
||||
},
|
||||
},
|
||||
)
|
||||
)
|
||||
recorder.handle_stream_event(
|
||||
AgentBackendStreamInternalEvent(
|
||||
run_id="run-1",
|
||||
data={
|
||||
"event_kind": "function_tool_call",
|
||||
"part": {
|
||||
"part_kind": "tool-call",
|
||||
"tool_name": "shell_run",
|
||||
"args": {"script": "lookup out"},
|
||||
},
|
||||
},
|
||||
)
|
||||
)
|
||||
recorder.handle_stream_event(
|
||||
AgentBackendStreamInternalEvent(
|
||||
run_id="run-1",
|
||||
data={
|
||||
"event_kind": "function_tool_result",
|
||||
"part": {
|
||||
"part_kind": "tool-return",
|
||||
"tool_name": "shell_run",
|
||||
"content": "out output",
|
||||
},
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
rows = sorted(fake_session.rows.values(), key=lambda row: row.position)
|
||||
assert len(rows) == 2
|
||||
assert rows[0].tool == "shell_run"
|
||||
assert rows[0].tool_input == '{"script": "lookup find"}'
|
||||
assert rows[0].observation == "find output"
|
||||
assert rows[1].tool == "shell_run"
|
||||
assert rows[1].tool_input == '{"script": "lookup out"}'
|
||||
assert rows[1].observation == "out output"
|
||||
|
||||
|
||||
def test_repeated_tool_calls_with_placeholder_call_id_and_reused_index_create_distinct_rows(monkeypatch):
|
||||
fake_session = _FakeDbSession()
|
||||
monkeypatch.setattr(app_runner_module.db, "session", fake_session)
|
||||
qm = _FakeQueueManager()
|
||||
recorder = app_runner_module._AgentProcessRecorder(
|
||||
dify_context=_dify_ctx(),
|
||||
message_id="msg-1",
|
||||
queue_manager=qm, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
for script, output in (("lookup find", "find output"), ("lookup out", "out output")):
|
||||
recorder.handle_stream_event(
|
||||
AgentBackendStreamInternalEvent(
|
||||
run_id="run-1",
|
||||
data={
|
||||
"event_kind": "function_tool_call",
|
||||
"index": 0,
|
||||
"part": {
|
||||
"part_kind": "tool-call",
|
||||
"tool_name": "shell_run",
|
||||
"tool_call_id": "None",
|
||||
"args": {"script": script},
|
||||
},
|
||||
},
|
||||
)
|
||||
)
|
||||
recorder.handle_stream_event(
|
||||
AgentBackendStreamInternalEvent(
|
||||
run_id="run-1",
|
||||
data={
|
||||
"event_kind": "function_tool_result",
|
||||
"part": {
|
||||
"part_kind": "tool-return",
|
||||
"tool_name": "shell_run",
|
||||
"tool_call_id": "None",
|
||||
"content": output,
|
||||
},
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
rows = sorted(fake_session.rows.values(), key=lambda row: row.position)
|
||||
assert len(rows) == 2
|
||||
assert rows[0].tool == "shell_run"
|
||||
assert rows[0].tool_input == '{"script": "lookup find"}'
|
||||
assert rows[0].observation == "find output"
|
||||
assert rows[1].tool == "shell_run"
|
||||
assert rows[1].tool_input == '{"script": "lookup out"}'
|
||||
assert rows[1].observation == "out output"
|
||||
|
||||
|
||||
def test_prior_session_snapshot_is_threaded_into_request():
|
||||
prior = CompositorSessionSnapshot(layers=[])
|
||||
client = FakeAgentBackendRunClient()
|
||||
@ -1088,6 +1214,19 @@ def test_failed_run_raises_agent_backend_error():
|
||||
assert store.saved == []
|
||||
|
||||
|
||||
def test_agent_backend_failure_to_exception_maps_rate_limit_reason():
|
||||
err = app_runner_module._agent_backend_failure_to_exception(
|
||||
AgentBackendRunFailedInternalEvent(
|
||||
run_id="run-1",
|
||||
error="quota exceeded",
|
||||
reason="InvokeRateLimitError",
|
||||
)
|
||||
)
|
||||
|
||||
assert isinstance(err, InvokeRateLimitError)
|
||||
assert str(err) == "quota exceeded"
|
||||
|
||||
|
||||
def test_stopped_task_cancels_agent_backend_run_and_skips_session_save():
|
||||
client = _RecordingFakeAgentBackendRunClient()
|
||||
store = _FakeSessionStore()
|
||||
|
||||
@ -3,10 +3,11 @@ from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
from core.app.apps.base_app_generate_response_converter import AppGenerateResponseConverter
|
||||
from core.app.entities.queue_entities import QueueErrorEvent
|
||||
from core.app.task_pipeline.based_generate_task_pipeline import BasedGenerateTaskPipeline
|
||||
from core.errors.error import QuotaExceededError
|
||||
from graphon.model_runtime.errors.invoke import InvokeAuthorizationError, InvokeError
|
||||
from graphon.model_runtime.errors.invoke import InvokeAuthorizationError, InvokeError, InvokeRateLimitError
|
||||
from models.enums import MessageStatus
|
||||
|
||||
|
||||
@ -68,6 +69,11 @@ class TestBasedGenerateTaskPipeline:
|
||||
assert error_response.task_id == "task-1"
|
||||
assert ping_response.task_id == "task-1"
|
||||
|
||||
def test_stream_converter_maps_invoke_rate_limit_error(self):
|
||||
data = AppGenerateResponseConverter._error_to_stream_response(InvokeRateLimitError("quota exceeded"))
|
||||
|
||||
assert data == {"code": "rate_limit_error", "status": 429, "message": "quota exceeded"}
|
||||
|
||||
def test_handle_output_moderation_when_flagged(self, pipeline):
|
||||
handler = Mock()
|
||||
handler.moderation_completion.return_value = ("filtered", True)
|
||||
|
||||
@ -22,6 +22,7 @@ def test_request_download_file_accepts_tool_file_reference() -> None:
|
||||
|
||||
assert payload.file.transfer_method == "tool_file"
|
||||
assert payload.file.reference == reference
|
||||
assert payload.for_external is True
|
||||
|
||||
|
||||
def test_request_download_file_accepts_remote_url() -> None:
|
||||
@ -42,6 +43,25 @@ def test_request_download_file_accepts_remote_url() -> None:
|
||||
assert payload.file.url == "https://example.com/report.pdf"
|
||||
|
||||
|
||||
def test_request_download_file_accepts_internal_download_request() -> None:
|
||||
reference = build_file_reference(record_id="tool-file-1")
|
||||
payload = RequestRequestDownloadFile.model_validate(
|
||||
{
|
||||
"tenant_id": "tenant-1",
|
||||
"user_id": "user-1",
|
||||
"user_from": "account",
|
||||
"invoke_from": "debugger",
|
||||
"file": {
|
||||
"transfer_method": "tool_file",
|
||||
"reference": reference,
|
||||
},
|
||||
"for_external": False,
|
||||
}
|
||||
)
|
||||
|
||||
assert payload.for_external is False
|
||||
|
||||
|
||||
def test_request_download_file_rejects_remote_url_without_url() -> None:
|
||||
with pytest.raises(ValidationError, match="url is required"):
|
||||
_ = RequestRequestDownloadFile.model_validate(
|
||||
|
||||
@ -60,6 +60,37 @@ def test_create_file_by_raw_stores_file_and_persists_record() -> None:
|
||||
session.refresh.assert_called_once_with(file_model)
|
||||
|
||||
|
||||
def test_create_file_by_raw_prefers_filename_extension_over_mimetype() -> None:
|
||||
manager = ToolFileManager()
|
||||
session = Mock()
|
||||
session.refresh.side_effect = lambda model: setattr(model, "id", "tf-docx")
|
||||
|
||||
def tool_file_factory(**kwargs):
|
||||
return SimpleNamespace(**kwargs)
|
||||
|
||||
with (
|
||||
patch("core.tools.tool_file_manager.storage") as storage,
|
||||
patch("core.tools.tool_file_manager.ToolFile", side_effect=tool_file_factory),
|
||||
patch("core.tools.tool_file_manager.uuid4", return_value=SimpleNamespace(hex="abc")),
|
||||
_patch_session_factory(session),
|
||||
):
|
||||
file_model = manager.create_file_by_raw(
|
||||
user_id="u1",
|
||||
tenant_id="t1",
|
||||
conversation_id="c1",
|
||||
file_binary=b"docx",
|
||||
mimetype="application/octet-stream",
|
||||
filename="report.docx",
|
||||
)
|
||||
|
||||
assert file_model.name == "report.docx"
|
||||
assert file_model.file_key == "tools/t1/abc.docx"
|
||||
storage.save.assert_called_once_with("tools/t1/abc.docx", b"docx")
|
||||
session.add.assert_called_once_with(file_model)
|
||||
session.commit.assert_called_once()
|
||||
session.refresh.assert_called_once_with(file_model)
|
||||
|
||||
|
||||
def test_create_file_by_url_downloads_and_persists_record() -> None:
|
||||
manager = ToolFileManager()
|
||||
response = Mock()
|
||||
@ -88,6 +119,32 @@ def test_create_file_by_url_downloads_and_persists_record() -> None:
|
||||
session.refresh.assert_called_once_with(file_model)
|
||||
|
||||
|
||||
def test_create_file_by_url_prefers_url_extension_over_mimetype() -> None:
|
||||
manager = ToolFileManager()
|
||||
response = Mock()
|
||||
response.content = b"docx"
|
||||
response.headers = {"Content-Type": "application/octet-stream"}
|
||||
response.raise_for_status.return_value = None
|
||||
session = Mock()
|
||||
|
||||
def tool_file_factory(**kwargs):
|
||||
return SimpleNamespace(**kwargs)
|
||||
|
||||
session.refresh.side_effect = lambda model: setattr(model, "id", "tf-docx")
|
||||
with (
|
||||
patch("core.tools.tool_file_manager.storage") as storage,
|
||||
patch("core.tools.tool_file_manager.ToolFile", side_effect=tool_file_factory),
|
||||
patch("core.tools.tool_file_manager.uuid4", return_value=SimpleNamespace(hex="urlabc")),
|
||||
_patch_session_factory(session),
|
||||
patch("core.tools.tool_file_manager.remote_fetcher.make_request", return_value=response),
|
||||
):
|
||||
file_model = manager.create_file_by_url("u1", "t1", "https://example.com/report.docx?download=1", "c1")
|
||||
|
||||
assert file_model.file_key == "tools/t1/urlabc.docx"
|
||||
assert file_model.name == "urlabc.docx"
|
||||
storage.save.assert_called_once_with("tools/t1/urlabc.docx", b"docx")
|
||||
|
||||
|
||||
def test_create_file_by_url_raises_on_timeout() -> None:
|
||||
manager = ToolFileManager()
|
||||
|
||||
|
||||
@ -7,9 +7,10 @@ from core.tools.entities.tool_entities import ToolInvokeMessage
|
||||
|
||||
|
||||
class _FakeToolFile:
|
||||
def __init__(self, mimetype: str):
|
||||
def __init__(self, mimetype: str, name: str | None):
|
||||
self.id = "fake-tool-file-id"
|
||||
self.mimetype = mimetype
|
||||
self.name = name or "fake-tool-file.bin"
|
||||
|
||||
|
||||
class _FakeToolFileManager:
|
||||
@ -38,7 +39,7 @@ class _FakeToolFileManager:
|
||||
"mimetype": mimetype,
|
||||
"filename": filename,
|
||||
}
|
||||
return _FakeToolFile(mimetype)
|
||||
return _FakeToolFile(mimetype, filename)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
@ -89,6 +90,29 @@ def test_transform_tool_invoke_messages_mimetype_key_present_but_none():
|
||||
assert o.meta["tool_file_id"] == "fake-tool-file-id"
|
||||
|
||||
|
||||
def test_transform_tool_invoke_messages_prefers_filename_extension_over_mimetype():
|
||||
msg = ToolInvokeMessage(
|
||||
type=ToolInvokeMessage.MessageType.BLOB,
|
||||
message=ToolInvokeMessage.BlobMessage(blob=b"docx"),
|
||||
meta={"mime_type": "application/octet-stream", "filename": "report.docx"},
|
||||
)
|
||||
|
||||
out = list(
|
||||
mt.ToolFileMessageTransformer.transform_tool_invoke_messages(
|
||||
messages=_gen([msg]),
|
||||
user_id="u1",
|
||||
tenant_id="t1",
|
||||
conversation_id="c1",
|
||||
)
|
||||
)
|
||||
|
||||
assert _FakeToolFileManager.last_call is not None
|
||||
assert _FakeToolFileManager.last_call["filename"] == "report.docx"
|
||||
assert len(out) == 1
|
||||
assert isinstance(out[0].message, ToolInvokeMessage.TextMessage)
|
||||
assert out[0].message.text.endswith(".docx")
|
||||
|
||||
|
||||
def test_transform_tool_invoke_messages_parses_existing_tool_file_link_meta():
|
||||
msg = ToolInvokeMessage(
|
||||
type=ToolInvokeMessage.MessageType.IMAGE_LINK,
|
||||
|
||||
@ -1,10 +1,12 @@
|
||||
from datetime import UTC, datetime
|
||||
from types import SimpleNamespace
|
||||
from typing import cast
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from agenton.compositor import CompositorSessionSnapshot
|
||||
from dify_agent.layers.ask_human import AskHumanToolResult
|
||||
from dify_agent.protocol import RunStartedEvent, RunSucceededEvent, RunSucceededEventData
|
||||
from dify_agent.protocol import PydanticAIStreamRunEvent, RunStartedEvent, RunSucceededEvent, RunSucceededEventData
|
||||
from pydantic_ai.messages import PartDeltaEvent, TextPartDelta
|
||||
|
||||
from clients.agent_backend import (
|
||||
AgentBackendRunEventAdapter,
|
||||
@ -190,6 +192,30 @@ class FileOutputBackendClient(FakeAgentBackendRunClient):
|
||||
)
|
||||
|
||||
|
||||
class AgentMessageDeltaBackendClient(FakeAgentBackendRunClient):
|
||||
def _events(self, run_id: str):
|
||||
created_at = datetime(2026, 1, 1, tzinfo=UTC)
|
||||
return (
|
||||
RunStartedEvent(id="1-0", run_id=run_id, created_at=created_at),
|
||||
PydanticAIStreamRunEvent(
|
||||
id="2-0",
|
||||
run_id=run_id,
|
||||
created_at=created_at,
|
||||
data=PartDeltaEvent(index=0, delta=TextPartDelta(content_delta="hello ")),
|
||||
agent_message_delta="hello ",
|
||||
),
|
||||
RunSucceededEvent(
|
||||
id="3-0",
|
||||
run_id=run_id,
|
||||
created_at=created_at,
|
||||
data=RunSucceededEventData(
|
||||
output={"text": "hello agent"},
|
||||
session_snapshot=CompositorSessionSnapshot(layers=[]),
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _node(
|
||||
*,
|
||||
scenario: FakeAgentBackendScenario = FakeAgentBackendScenario.SUCCESS,
|
||||
@ -277,6 +303,19 @@ def test_agent_node_run_maps_successful_agent_backend_run_to_node_result():
|
||||
assert layers["llm"]["config"]["credentials"] == "[REDACTED]"
|
||||
|
||||
|
||||
def test_agent_node_run_ignores_agent_message_delta_until_terminal_result():
|
||||
events = list(_node(agent_backend_client=AgentMessageDeltaBackendClient())._run())
|
||||
|
||||
assert len(events) == 1
|
||||
result = cast(StreamCompletedEvent, events[0]).node_run_result
|
||||
assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED
|
||||
assert result.outputs == {"text": "hello agent"}
|
||||
agent_backend = result.metadata[WorkflowNodeExecutionMetadataKey.AGENT_LOG]["agent_backend"]
|
||||
assert agent_backend["status"] == "succeeded"
|
||||
assert agent_backend["agent_message_delta_count"] == 1
|
||||
assert agent_backend["agent_message_delta_length"] == len("hello ")
|
||||
|
||||
|
||||
def test_agent_node_run_normalizes_declared_file_output_with_canonical_mapping():
|
||||
tool_reference = build_file_reference(record_id="tool-file-1")
|
||||
with patch(
|
||||
|
||||
@ -57,6 +57,17 @@ def test_reback_resolves_tenant_tool_file_to_file():
|
||||
assert file.extension == ".png"
|
||||
|
||||
|
||||
def test_reback_prefers_filename_extension_over_mimetype():
|
||||
tf = _seed(mimetype="application/octet-stream", name="report.docx", size=99)
|
||||
file = reback_tool_file_output(tenant_id=TENANT, tool_file_id=tf)
|
||||
|
||||
assert file is not None
|
||||
assert file.filename == "report.docx"
|
||||
assert file.mime_type == "application/octet-stream"
|
||||
assert file.extension == ".docx"
|
||||
assert file.type == FileType.CUSTOM
|
||||
|
||||
|
||||
def test_reback_other_tenant_returns_none():
|
||||
tf = _seed()
|
||||
assert reback_tool_file_output(tenant_id="33333333-3333-3333-3333-333333333333", tool_file_id=tf) is None
|
||||
|
||||
@ -165,6 +165,20 @@ def test_build_from_mapping_accepts_opaque_related_id_for_tool_file(mock_tool_fi
|
||||
assert file.storage_key == "tool_file.pdf"
|
||||
|
||||
|
||||
def test_build_from_mapping_prefers_tool_filename_extension_over_mimetype(mock_tool_file):
|
||||
mock_tool_file.name = "report.docx"
|
||||
mock_tool_file.file_key = "tools/test_tenant_id/file.bin"
|
||||
mock_tool_file.mimetype = "application/octet-stream"
|
||||
mapping = tool_file_mapping(file_type="document")
|
||||
|
||||
file = build_from_mapping(mapping=mapping, tenant_id=TEST_TENANT_ID)
|
||||
|
||||
assert file.extension == ".docx"
|
||||
assert file.filename == "report.docx"
|
||||
assert file.mime_type == "application/octet-stream"
|
||||
assert file.storage_key == "tools/test_tenant_id/file.bin"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("file_type", "should_pass", "expected_error"),
|
||||
[
|
||||
@ -213,6 +227,25 @@ def test_build_from_remote_url(mock_http_head):
|
||||
assert file.size == 2048
|
||||
|
||||
|
||||
def test_build_from_remote_url_prefers_filename_extension_over_mimetype():
|
||||
mapping = {
|
||||
"transfer_method": "remote_url",
|
||||
"url": TEST_REMOTE_URL,
|
||||
"type": "document",
|
||||
}
|
||||
|
||||
with patch(
|
||||
"factories.file_factory.builders.get_remote_file_info",
|
||||
return_value=("application/octet-stream", "report.docx", 99),
|
||||
):
|
||||
file = build_from_mapping(mapping=mapping, tenant_id=TEST_TENANT_ID)
|
||||
|
||||
assert file.filename == "report.docx"
|
||||
assert file.extension == ".docx"
|
||||
assert file.mime_type == "application/octet-stream"
|
||||
assert file.size == 99
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("file_type", "should_pass", "expected_error"),
|
||||
[
|
||||
|
||||
@ -1,4 +1,6 @@
|
||||
from fields.message_fields import ExploreMessageListItem, MessageListItem
|
||||
from decimal import Decimal
|
||||
|
||||
from fields.message_fields import ExploreMessageListItem, MessageListItem, WebMessageListItem
|
||||
|
||||
|
||||
def _base_kwargs():
|
||||
@ -34,3 +36,33 @@ class TestExploreMessageListItem:
|
||||
# Guard the public service-API contract: the base item must not leak metadata.
|
||||
payload = MessageListItem(**_base_kwargs()).model_dump(mode="json")
|
||||
assert "metadata" not in payload
|
||||
|
||||
def test_message_list_item_exposes_usage_fields(self):
|
||||
payload = MessageListItem(
|
||||
**_base_kwargs(),
|
||||
message_tokens=7,
|
||||
answer_tokens=11,
|
||||
provider_response_latency=1.25,
|
||||
total_price=Decimal("0.0001234"),
|
||||
currency="USD",
|
||||
).model_dump(mode="json")
|
||||
|
||||
assert payload["message_tokens"] == 7
|
||||
assert payload["answer_tokens"] == 11
|
||||
assert payload["total_tokens"] == 18
|
||||
assert payload["provider_response_latency"] == 1.25
|
||||
assert payload["total_price"] == "0.0001234"
|
||||
assert payload["currency"] == "USD"
|
||||
|
||||
def test_web_message_list_item_exposes_usage_and_metadata(self):
|
||||
payload = WebMessageListItem(
|
||||
**_base_kwargs(),
|
||||
metadata={"usage": {"total_tokens": 18}},
|
||||
message_tokens=7,
|
||||
answer_tokens=11,
|
||||
).model_dump(mode="json")
|
||||
|
||||
assert payload["metadata"] == {"usage": {"total_tokens": 18}}
|
||||
assert payload["message_tokens"] == 7
|
||||
assert payload["answer_tokens"] == 11
|
||||
assert payload["total_tokens"] == 18
|
||||
|
||||
@ -59,6 +59,31 @@ def test_request_download_url_builds_file_under_bound_scope(
|
||||
assert result.download_url == "https://files.example.com/x"
|
||||
|
||||
|
||||
def test_request_download_url_supports_internal_download_urls() -> None:
|
||||
fake_file = MagicMock(filename="report.pdf", mime_type="application/pdf", size=123)
|
||||
service = FileRequestService(access_controller=MagicMock())
|
||||
|
||||
with (
|
||||
patch("services.file_request_service.bind_file_access_scope", return_value=nullcontext()),
|
||||
patch.object(service, "_build_file", return_value=fake_file),
|
||||
patch(
|
||||
"services.file_request_service.file_helpers.resolve_file_url",
|
||||
return_value="http://internal-files/report.pdf",
|
||||
) as resolve_file_url,
|
||||
):
|
||||
result = service.request_download_url(
|
||||
tenant_id="tenant-1",
|
||||
user_id="user-1",
|
||||
user_from="account",
|
||||
invoke_from="debugger",
|
||||
file_mapping={"transfer_method": "tool_file", "reference": "dify-file-ref:tool-file-1"},
|
||||
for_external=False,
|
||||
)
|
||||
|
||||
resolve_file_url.assert_called_once_with(fake_file, for_external=False)
|
||||
assert result.download_url == "http://internal-files/report.pdf"
|
||||
|
||||
|
||||
def test_request_download_url_rejects_unsupported_files() -> None:
|
||||
service = FileRequestService(access_controller=MagicMock())
|
||||
|
||||
|
||||
37
api/uv.lock
generated
37
api/uv.lock
generated
@ -322,16 +322,15 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "anyio"
|
||||
version = "4.11.0"
|
||||
version = "4.14.1"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "idna" },
|
||||
{ name = "sniffio" },
|
||||
{ name = "typing-extensions" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/c6/78/7d432127c41b50bccba979505f272c16cbcadcc33645d5fa3a738110ae75/anyio-4.11.0.tar.gz", hash = "sha256:82a8d0b81e318cc5ce71a5f1f8b5c4e63619620b63141ef8c995fa0db95a57c4", size = 219094, upload-time = "2025-09-23T09:19:12.58Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/3b/72/5562aabb8dd7181e8e860622a38bea08d17842b99ecd4c91f84ac95251b0/anyio-4.14.1.tar.gz", hash = "sha256:8d648a3544c1a700e3ff78615cd679e4c5c3f149904287e73687b2596963629e", size = 254831, upload-time = "2026-06-24T20:56:06.017Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/15/b3/9b1a8074496371342ec1e796a96f99c82c945a339cd81a8e73de28b4cf9e/anyio-4.11.0-py3-none-any.whl", hash = "sha256:0287e96f4d26d4149305414d4e3bc32f0dcd0862365a4bddea19d7a1ec38c4fc", size = 109097, upload-time = "2025-09-23T09:19:10.601Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/b0/7b/90df4a0a816d98d6ea26f559d87836d494a2cf1fcf063be67df50a7bcc30/anyio-4.14.1-py3-none-any.whl", hash = "sha256:4e5533c5b8ff0a24f5d7a176cbe6877129cd183893f66b537f8f227d10527d72", size = 124875, upload-time = "2026-06-24T20:56:04.413Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@ -1284,7 +1283,9 @@ name = "dify-agent"
|
||||
version = "0.1.0"
|
||||
source = { editable = "../dify-agent" }
|
||||
dependencies = [
|
||||
{ name = "anyio" },
|
||||
{ name = "httpx" },
|
||||
{ name = "httpx2" },
|
||||
{ name = "pydantic" },
|
||||
{ name = "pydantic-ai-slim" },
|
||||
{ name = "typer" },
|
||||
@ -1293,10 +1294,14 @@ dependencies = [
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "aiosqlite", marker = "extra == 'shellctl-server'", specifier = ">=0.21.0,<1.0.0" },
|
||||
{ name = "anyio", specifier = ">=4.12.1,<5.0.0" },
|
||||
{ name = "fastapi", marker = "extra == 'server'", specifier = "==0.136.0" },
|
||||
{ name = "fastapi", marker = "extra == 'shellctl-server'", specifier = "==0.136.0" },
|
||||
{ name = "graphon", marker = "extra == 'server'", specifier = "==0.5.2" },
|
||||
{ name = "grpclib", extras = ["protobuf"], marker = "extra == 'grpc'", specifier = ">=0.4.9,<0.5.0" },
|
||||
{ name = "httpx", specifier = "==0.28.1" },
|
||||
{ name = "httpx2", specifier = ">=2.5.0,<3.0.0" },
|
||||
{ name = "jsonschema", marker = "extra == 'server'", specifier = ">=4.23.0,<5.0.0" },
|
||||
{ name = "jwcrypto", marker = "extra == 'server'", specifier = ">=1.5.6,<2" },
|
||||
{ name = "logfire", extras = ["fastapi", "httpx", "redis"], marker = "extra == 'server'", specifier = ">=4.37.0,<5.0.0" },
|
||||
@ -1306,12 +1311,13 @@ requires-dist = [
|
||||
{ name = "pydantic-ai-slim", extras = ["anthropic", "google", "openai"], marker = "extra == 'server'", specifier = ">=1.85.1,<2.0.0" },
|
||||
{ name = "pydantic-settings", marker = "extra == 'server'", specifier = ">=2.12.0,<3.0.0" },
|
||||
{ name = "redis", marker = "extra == 'server'", specifier = ">=7.4.0,<8.0.0" },
|
||||
{ name = "shell-session-manager", marker = "extra == 'server'", specifier = "==2.4.0" },
|
||||
{ name = "sqlmodel", marker = "extra == 'shellctl-server'", specifier = ">=0.0.24,<0.1.0" },
|
||||
{ name = "typer", specifier = ">=0.16.1,<0.17" },
|
||||
{ name = "typing-extensions", specifier = ">=4.12.2,<5.0.0" },
|
||||
{ name = "uvicorn", extras = ["standard"], marker = "extra == 'server'", specifier = "==0.46.0" },
|
||||
{ name = "uvicorn", extras = ["standard"], marker = "extra == 'shellctl-server'", specifier = "==0.46.0" },
|
||||
]
|
||||
provides-extras = ["grpc", "server"]
|
||||
provides-extras = ["grpc", "server", "shellctl-server"]
|
||||
|
||||
[package.metadata.requires-dev]
|
||||
dev = [
|
||||
@ -3283,15 +3289,15 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "httpcore2"
|
||||
version = "2.3.0"
|
||||
version = "2.5.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "h11" },
|
||||
{ name = "truststore" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/e6/34/18f1c596e677962f040284246f393b10a1f8ce440b3a7e69c637d0f1c7ad/httpcore2-2.3.0.tar.gz", hash = "sha256:07327e251560960eea8e969d92d4c6a325feb13cca39e25340731336c3baf924", size = 64300, upload-time = "2026-06-01T13:15:02.998Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/47/06/5c12df521b5322fb1114a83d46911b2fbcb8855ddb3a635f11c01a214af5/httpcore2-2.5.0.tar.gz", hash = "sha256:88aa170137c17328d5ac44234f9fd10706466d5fb347f3edac4d39b91137b09d", size = 64808, upload-time = "2026-06-25T14:16:56.472Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/c2/dd/3357218c69360d1cecc196c230c9a1d5c9afd5dba362056e23e60a5e64e5/httpcore2-2.3.0-py3-none-any.whl", hash = "sha256:477e9e334f74e5240dcac002e890580f36a57d40ff0fb14cc9655731d23b8415", size = 80024, upload-time = "2026-06-01T13:15:00.001Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/c9/a1/7564199d1a8728fe737b0a72e5b3f8d92dfe085a74ddf7cdd83bce5f206d/httpcore2-2.5.0-py3-none-any.whl", hash = "sha256:5ce35188de461d31e8d000bfb8ef8bf22c6c16587a211e5571deaa5e9bdf842a", size = 80330, upload-time = "2026-06-25T14:16:53.634Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@ -3355,17 +3361,18 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "httpx2"
|
||||
version = "2.3.0"
|
||||
version = "2.5.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "anyio" },
|
||||
{ name = "httpcore2" },
|
||||
{ name = "idna" },
|
||||
{ name = "truststore" },
|
||||
{ name = "typing-extensions" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/9f/9a/cca0b9145f13d8ae34b885ae28d403a1469a433abc78e0f94f4ce94e650b/httpx2-2.3.0.tar.gz", hash = "sha256:227e7c41d95a76d4077a52640564132777215fc3394e07b66a3116c33d668fa9", size = 81115, upload-time = "2026-06-01T13:15:04.324Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/d0/e2/b5dedc0cf35aa65de5f541ccd30d2bc1fd7f1d43c9ab09f8ed9a7342317b/httpx2-2.5.0.tar.gz", hash = "sha256:e2df9cb4611021527ff8a675b1c320b610a2ec397acc8d6fe6e91df2d9b33c29", size = 83121, upload-time = "2026-06-25T14:16:57.491Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/87/ce/ae2911859847f9ba1d6b23027e53481cbeb50b93234f355a968d300ca2cb/httpx2-2.3.0-py3-none-any.whl", hash = "sha256:6f393663bdf6dbe7fe90118e3eb5b2bd024a675cae0390ac08cec9198812d8b7", size = 74538, upload-time = "2026-06-01T13:15:01.566Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/31/22/859d8252dad9bc9adee34b52e62cde621ece07b042ccb2ab4da1be46695f/httpx2-2.5.0-py3-none-any.whl", hash = "sha256:3d2d4d9cf4b61f1a1f46a95947cfdb47e80cb56a2f91c6256ac8f58e4891df41", size = 76652, upload-time = "2026-06-25T14:16:55.23Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@ -3423,11 +3430,11 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "idna"
|
||||
version = "3.11"
|
||||
version = "3.18"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/6f/6d/0703ccc57f3a7233505399edb88de3cbd678da106337b9fcde432b65ed60/idna-3.11.tar.gz", hash = "sha256:795dafcc9c04ed0c1fb032c2aa73654d8e8c5023a7df64a53f39190ada629902", size = 194582, upload-time = "2025-10-12T14:55:20.501Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/cd/63/9496c57188a2ee585e0f1db071d75089a11e98aa86eb99d9d7618fc1edce/idna-3.18.tar.gz", hash = "sha256:ffb385a7e039654cef1ab9ef32c6fafe283c0c0467bba1d9029738ce4a14a848", size = 196711, upload-time = "2026-06-02T14:34:07.794Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/0e/61/66938bbb5fc52dbdf84594873d5b51fb1f7c7794e9c0f5bd885f30bc507b/idna-3.11-py3-none-any.whl", hash = "sha256:771a87f49d9defaf64091e6e6fe9c18d4833f140bd19464795bc32d966ca37ea", size = 71008, upload-time = "2025-10-12T14:55:18.883Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/1e/5e/d4e9f1a599fb8e573b7b87160658329fbf28d19eac2718f51fc3def3aa5a/idna-3.18-py3-none-any.whl", hash = "sha256:7f952cbe720b688055e3f87de14f5c3e5fdaa8bc3928985c4077ca689de849a2", size = 65455, upload-time = "2026-06-02T14:34:06.319Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
@ -20,6 +20,12 @@ DIFY_AGENT_PLUGIN_DAEMON_URL=http://localhost:5002
|
||||
# API key sent to the Dify plugin daemon.
|
||||
DIFY_AGENT_PLUGIN_DAEMON_API_KEY=
|
||||
|
||||
# Dify API inner endpoints
|
||||
# Base URL for Dify API inner endpoints used by Agent Stub config/file/drive requests.
|
||||
DIFY_AGENT_INNER_API_URL=http://localhost:5001
|
||||
# Must match API/worker INNER_API_KEY_FOR_PLUGIN, not the generic INNER_API_KEY.
|
||||
DIFY_AGENT_INNER_API_KEY=
|
||||
|
||||
# Shell layer
|
||||
# Base URL for the shellctl server used by the dify.shell layer. Leave empty to disable shell layer use.
|
||||
DIFY_AGENT_SHELLCTL_ENTRYPOINT=
|
||||
@ -30,12 +36,14 @@ DIFY_AGENT_SHELLCTL_AUTH_TOKEN=
|
||||
# Public Agent Stub URL reachable from shellctl-managed remote machines.
|
||||
# Use http(s)://.../agent-stub for HTTP or grpc://host:port for gRPC.
|
||||
# Leave empty to avoid injecting DIFY_AGENT_STUB_* into shell.run jobs.
|
||||
DIFY_AGENT_STUB_URL=
|
||||
# Optional bind override used only when DIFY_AGENT_STUB_URL uses grpc://.
|
||||
DIFY_AGENT_STUB_API_BASE_URL=http://localhost:5050/agent-stub
|
||||
# Optional bind override used only when DIFY_AGENT_STUB_API_BASE_URL uses grpc://.
|
||||
DIFY_AGENT_STUB_GRPC_BIND_ADDRESS=
|
||||
# Server-wide root secret used to derive Agent Stub JWE keys.
|
||||
# Required when DIFY_AGENT_STUB_URL is set; must be unpadded base64url for 32 bytes.
|
||||
DIFY_AGENT_SERVER_SECRET_KEY=
|
||||
# This is security-sensitive: it derives the JWE encryption key for Agent Stub bearer tokens.
|
||||
# Replace this development default in production.
|
||||
# Generate one with: python -c 'import secrets; print(secrets.token_urlsafe(32))'
|
||||
DIFY_AGENT_SERVER_SECRET_KEY=MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY
|
||||
|
||||
# Shared plugin-daemon HTTP client timeouts and limits.
|
||||
# Plugin-daemon HTTP connect timeout in seconds.
|
||||
|
||||
@ -6,8 +6,8 @@
|
||||
# cd /app/api && .venv/bin/uvicorn dify_agent.server.app:app --host 0.0.0.0 --port 5050
|
||||
#
|
||||
# Unlike the dify-api image (which only installs the base `dify-agent`
|
||||
# dependency), this image installs the `[server]` extra, so jwcrypto,
|
||||
# shell-session-manager, fastapi, uvicorn, etc. are present and the server can
|
||||
# dependency), this image installs the `[server]` extra, so jwcrypto, fastapi,
|
||||
# uvicorn, etc. are present and the server can
|
||||
# actually start. dify-api is intentionally left lean.
|
||||
|
||||
# base image
|
||||
|
||||
@ -12,19 +12,15 @@ FROM python:3.12-slim-bookworm AS base
|
||||
ARG NODE_VERSION=22.22.1
|
||||
ARG PNPM_VERSION=11.9.0
|
||||
ARG UV_VERSION=0.8.9
|
||||
ARG DIFY_AGENT_TOOL_SPEC=.[grpc]
|
||||
ARG SHELL_SESSION_MANAGER_TOOL_SPEC=shell-session-manager==2.4.0
|
||||
ARG DIFY_AGENT_TOOL_SPEC=.[grpc,shellctl-server]
|
||||
|
||||
ENV PYTHONDONTWRITEBYTECODE=1 \
|
||||
PYTHONUNBUFFERED=1 \
|
||||
PIP_NO_CACHE_DIR=1 \
|
||||
DIFY_AGENT_STUB_DRIVE_BASE=/mnt/drive \
|
||||
UV_TOOL_DIR=/opt/dify-agent-tools/envs \
|
||||
UV_TOOL_BIN_DIR=/opt/dify-agent-tools/bin
|
||||
ENV PATH="${UV_TOOL_BIN_DIR}:${PATH}"
|
||||
PIP_NO_CACHE_DIR=1
|
||||
|
||||
RUN apt-get update \
|
||||
&& apt-get install -y --no-install-recommends \
|
||||
bash \
|
||||
ca-certificates \
|
||||
curl \
|
||||
file \
|
||||
@ -62,24 +58,25 @@ WORKDIR /opt/dify-agent
|
||||
FROM base AS tools
|
||||
|
||||
ARG DIFY_AGENT_TOOL_SPEC
|
||||
ARG SHELL_SESSION_MANAGER_TOOL_SPEC
|
||||
|
||||
COPY pyproject.toml uv.lock README.md ./
|
||||
COPY src ./src
|
||||
|
||||
RUN uv export --frozen --no-dev --all-extras --no-emit-project --no-hashes \
|
||||
> /tmp/dify-agent-constraints.txt \
|
||||
&& uv tool install --force --python /usr/local/bin/python --no-python-downloads \
|
||||
&& UV_TOOL_DIR=/opt/dify-agent-tools/envs \
|
||||
UV_TOOL_BIN_DIR=/opt/dify-agent-tools/bin \
|
||||
uv tool install --force --python /usr/local/bin/python --no-python-downloads \
|
||||
--constraints /tmp/dify-agent-constraints.txt --link-mode=copy "${DIFY_AGENT_TOOL_SPEC}" \
|
||||
&& uv tool install --force --python /usr/local/bin/python --no-python-downloads \
|
||||
--constraints /tmp/dify-agent-constraints.txt --link-mode=copy "${SHELL_SESSION_MANAGER_TOOL_SPEC}" \
|
||||
&& rm -f /tmp/dify-agent-constraints.txt
|
||||
|
||||
|
||||
FROM base AS production
|
||||
|
||||
COPY --from=tools ${UV_TOOL_DIR} ${UV_TOOL_DIR}
|
||||
COPY --from=tools ${UV_TOOL_BIN_DIR} ${UV_TOOL_BIN_DIR}
|
||||
ENV PATH="/opt/dify-agent-tools/bin:${PATH}"
|
||||
|
||||
COPY --from=tools /opt/dify-agent-tools/envs /opt/dify-agent-tools/envs
|
||||
COPY --from=tools /opt/dify-agent-tools/bin /opt/dify-agent-tools/bin
|
||||
|
||||
RUN useradd --create-home --shell /bin/sh dify \
|
||||
&& mkdir -p /mnt/drive \
|
||||
|
||||
@ -75,8 +75,9 @@ See `.example.env` for the full server settings template.
|
||||
|
||||
If you plan to run `dify.shell`, also configure `DIFY_AGENT_SHELLCTL_ENTRYPOINT`
|
||||
and, when shell jobs need to call back with the `dify-agent` command, set
|
||||
`DIFY_AGENT_STUB_API_BASE_URL` plus a 32-byte base64url
|
||||
`DIFY_AGENT_SERVER_SECRET_KEY` as documented in `.example.env`.
|
||||
`DIFY_AGENT_STUB_API_BASE_URL`. The supplied default configs include a
|
||||
development `DIFY_AGENT_SERVER_SECRET_KEY`, but production deployments should
|
||||
override it with a unique 32-byte base64url value as documented in `.example.env`.
|
||||
|
||||
## Start the Dify Agent server
|
||||
|
||||
|
||||
@ -42,7 +42,7 @@ also reads `.env` and `dify-agent/.env` when present.
|
||||
| `DIFY_AGENT_SHELLCTL_AUTH_TOKEN` | empty | Optional bearer token sent to the shellctl server. |
|
||||
| `DIFY_AGENT_STUB_API_BASE_URL` | empty | Public Agent Stub API base URL reachable from shellctl-managed remote machines. HTTP may be the service root or `/agent-stub`; gRPC must be `grpc://host:port`. Enables `DIFY_AGENT_STUB_*` env injection for user `shell.run` jobs. |
|
||||
| `DIFY_AGENT_STUB_GRPC_BIND_ADDRESS` | empty | Optional `host:port` bind override used only when `DIFY_AGENT_STUB_API_BASE_URL` uses `grpc://`. |
|
||||
| `DIFY_AGENT_SERVER_SECRET_KEY` | empty | Server-wide root secret used to derive Agent Stub JWE keys; required when `DIFY_AGENT_STUB_API_BASE_URL` is set and must be unpadded base64url for 32 bytes. |
|
||||
| `DIFY_AGENT_SERVER_SECRET_KEY` | empty | Security-sensitive server-wide root secret used to derive the JWE encryption key for Agent Stub bearer tokens; required when `DIFY_AGENT_STUB_API_BASE_URL` is set. The supplied default config uses a development value; set a unique unpadded base64url 32-byte secret in production. |
|
||||
| `DIFY_AGENT_PLUGIN_DAEMON_CONNECT_TIMEOUT` | `10` | Plugin-daemon HTTP connect timeout in seconds. |
|
||||
| `DIFY_AGENT_PLUGIN_DAEMON_READ_TIMEOUT` | `600` | Plugin-daemon HTTP read timeout in seconds. |
|
||||
| `DIFY_AGENT_PLUGIN_DAEMON_WRITE_TIMEOUT` | `30` | Plugin-daemon HTTP write timeout in seconds. |
|
||||
@ -64,9 +64,11 @@ DIFY_AGENT_INNER_API_URL=http://localhost:5001
|
||||
DIFY_AGENT_INNER_API_KEY=replace-with-dify-inner-api-key-for-plugin
|
||||
DIFY_AGENT_SHELLCTL_ENTRYPOINT=http://127.0.0.1:5004
|
||||
DIFY_AGENT_SHELLCTL_AUTH_TOKEN=replace-with-shellctl-token
|
||||
# Generate with: python -c 'import base64, secrets; print(base64.urlsafe_b64encode(secrets.token_bytes(32)).rstrip(b"=").decode())'
|
||||
DIFY_AGENT_STUB_API_BASE_URL=https://agent.example.com/agent-stub
|
||||
DIFY_AGENT_SERVER_SECRET_KEY=replace-with-base64url-32-byte-secret
|
||||
# This is security-sensitive: it derives the JWE encryption key for Agent Stub bearer tokens.
|
||||
# Replace this development default in production.
|
||||
# Generate one with: python -c 'import secrets; print(secrets.token_urlsafe(32))'
|
||||
DIFY_AGENT_SERVER_SECRET_KEY=MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY
|
||||
```
|
||||
|
||||
Run records and event streams use the same retention. Status writes refresh the
|
||||
|
||||
@ -56,18 +56,22 @@ with `dify-agent ...`, also enable the Agent Stub:
|
||||
|
||||
```env
|
||||
DIFY_AGENT_STUB_API_BASE_URL=https://agent.example.com/agent-stub
|
||||
DIFY_AGENT_SERVER_SECRET_KEY=replace-with-base64url-32-byte-secret
|
||||
# This is security-sensitive: it derives the JWE encryption key for Agent Stub bearer tokens.
|
||||
# Replace this development default in production.
|
||||
# Generate one with: python -c 'import secrets; print(secrets.token_urlsafe(32))'
|
||||
DIFY_AGENT_SERVER_SECRET_KEY=MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY
|
||||
```
|
||||
|
||||
HTTP `DIFY_AGENT_STUB_API_BASE_URL` may be either the service root or the
|
||||
explicit `/agent-stub` API root; the server normalizes the service root to
|
||||
`/agent-stub`. Other HTTP paths are rejected at startup.
|
||||
|
||||
`DIFY_AGENT_SERVER_SECRET_KEY` must be unpadded base64url text for exactly 32
|
||||
decoded bytes. One way to generate it is:
|
||||
The supplied Docker and `.example.env` configs use a development
|
||||
`DIFY_AGENT_SERVER_SECRET_KEY`. Override it in production with unpadded base64url
|
||||
text for exactly 32 decoded bytes. One way to generate it is:
|
||||
|
||||
```bash
|
||||
python -c 'import base64, secrets; print(base64.urlsafe_b64encode(secrets.token_bytes(32)).rstrip(b"=").decode())'
|
||||
python -c 'import secrets; print(secrets.token_urlsafe(32))'
|
||||
```
|
||||
|
||||
## Client request shape
|
||||
@ -230,12 +234,11 @@ The provided `docker/local-sandbox/Dockerfile` installs:
|
||||
- `tmux`, required by `shellctl` to manage shell jobs;
|
||||
- common shell workspace tools: `git`, `openssh-client`, `jq`, `ripgrep`,
|
||||
`unzip`, `zip`, `file`, `procps`, and `less`;
|
||||
- `shell-session-manager==2.3.1` as a standalone uv tool, which provides the
|
||||
`shellctl` CLI/server;
|
||||
- `dify-agent[grpc,shellctl-server]` as a standalone uv tool, which provides
|
||||
both the Agent Stub client CLI and the built-in `shellctl` CLI/server;
|
||||
- `uv`, so uv shebang scripts with PEP 723 metadata can run inside the shell
|
||||
workspace and Python CLI tools can be installed with isolated tool
|
||||
environments;
|
||||
- `node==22.22.1` and `pnpm==11.9.0`, so JavaScript and TypeScript tooling can
|
||||
run inside the shell workspace without per-job installation;
|
||||
- the `dify-agent[grpc]` Agent Stub client CLI as a standalone uv tool;
|
||||
- a non-root default user named `dify`.
|
||||
|
||||
@ -36,6 +36,7 @@ message FileMapping {
|
||||
|
||||
message FileDownloadRequest {
|
||||
FileMapping file = 1;
|
||||
optional bool for_external = 2;
|
||||
}
|
||||
|
||||
message FileDownloadResponse {
|
||||
|
||||
@ -5,7 +5,9 @@ description = "Add your description here"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12,<4.0"
|
||||
dependencies = [
|
||||
"anyio>=4.12.1,<5.0.0",
|
||||
"httpx==0.28.1",
|
||||
"httpx2>=2.5.0,<3.0.0",
|
||||
"pydantic>=2.12.5,<2.13",
|
||||
"pydantic-ai-slim>=1.102.0,<2.0.0",
|
||||
"typer>=0.16.1,<0.17",
|
||||
@ -15,6 +17,9 @@ dependencies = [
|
||||
[project.scripts]
|
||||
dify-agent = "dify_agent.agent_stub.cli.main:main"
|
||||
dify-agent-stub-server = "dify_agent.agent_stub.server.cli:main"
|
||||
shellctl = "shellctl.cli:main"
|
||||
shellctl-sanitize-pty = "shellctl_runtime.sanitize:main"
|
||||
shellctl-runner-exit = "shellctl_runtime.runner_exit:main"
|
||||
|
||||
[project.optional-dependencies]
|
||||
grpc = ["grpclib[protobuf]>=0.4.9,<0.5.0", "protobuf>=6.33.5,<7.0.0"]
|
||||
@ -27,13 +32,18 @@ server = [
|
||||
"pydantic-ai-slim[anthropic,google,openai]>=1.85.1,<2.0.0",
|
||||
"pydantic-settings>=2.12.0,<3.0.0",
|
||||
"redis>=7.4.0,<8.0.0",
|
||||
"shell-session-manager==2.4.0",
|
||||
"uvicorn[standard]==0.46.0",
|
||||
]
|
||||
shellctl-server = [
|
||||
"aiosqlite>=0.21.0,<1.0.0",
|
||||
"fastapi==0.136.0",
|
||||
"sqlmodel>=0.0.24,<0.1.0",
|
||||
"uvicorn[standard]==0.46.0",
|
||||
]
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
where = ["src"]
|
||||
include = ["agenton*", "agenton_collections*", "dify_agent*"]
|
||||
include = ["agenton*", "agenton_collections*", "dify_agent*", "shellctl*", "shellctl_runtime*"]
|
||||
|
||||
[tool.pyright]
|
||||
include = ["src", "examples", "tests"]
|
||||
|
||||
@ -205,6 +205,7 @@ class DifyStreamedResponse(StreamedResponse):
|
||||
|
||||
@override
|
||||
async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
|
||||
chunk_sequence = 0
|
||||
async for chunk in self.chunks:
|
||||
if chunk.delta.usage is not None:
|
||||
self._usage: RequestUsage = _map_usage(chunk.delta.usage)
|
||||
@ -216,8 +217,10 @@ class DifyStreamedResponse(StreamedResponse):
|
||||
chunk,
|
||||
self.provider_name_value,
|
||||
self._embedded_thinking_parser,
|
||||
chunk_sequence,
|
||||
):
|
||||
yield event
|
||||
chunk_sequence += 1
|
||||
|
||||
for event in self._embedded_thinking_parser.flush(self._parts_manager, self.provider_name_value):
|
||||
yield event
|
||||
@ -551,11 +554,21 @@ def _normalize_finish_reason(finish_reason: str) -> FinishReason:
|
||||
return "error"
|
||||
|
||||
|
||||
def _normalize_tool_call_id(tool_call_id: str | None) -> str | None:
|
||||
if tool_call_id is None:
|
||||
return None
|
||||
normalized = tool_call_id.strip()
|
||||
if not normalized or normalized.lower() in {"none", "null"}:
|
||||
return None
|
||||
return normalized
|
||||
|
||||
|
||||
def _chunk_to_stream_events(
|
||||
parts_manager: ModelResponsePartsManager,
|
||||
chunk: LLMResultChunk,
|
||||
provider_name: str,
|
||||
embedded_thinking_parser: "_EmbeddedThinkingParser",
|
||||
chunk_sequence: int,
|
||||
) -> list[ModelResponseStreamEvent]:
|
||||
events: list[ModelResponseStreamEvent] = []
|
||||
message = chunk.delta.message
|
||||
@ -571,13 +584,14 @@ def _chunk_to_stream_events(
|
||||
events.append(parts_manager.handle_part(vendor_part_id=None, part=part))
|
||||
|
||||
for index, tool_call in enumerate(message.tool_calls):
|
||||
vendor_id = tool_call.id or f"chunk-{chunk.delta.index}-tool-{index}"
|
||||
tool_call_id = _normalize_tool_call_id(tool_call.id)
|
||||
vendor_id = tool_call_id or f"chunk-{chunk_sequence}-tool-{index}"
|
||||
events.append(
|
||||
parts_manager.handle_tool_call_part(
|
||||
vendor_part_id=vendor_id,
|
||||
tool_name=tool_call.function.name,
|
||||
args=tool_call.function.arguments,
|
||||
tool_call_id=tool_call.id,
|
||||
tool_call_id=tool_call_id or vendor_id,
|
||||
provider_name=provider_name,
|
||||
)
|
||||
)
|
||||
|
||||
@ -1,6 +1,6 @@
|
||||
"""Shellctl-backed shell provider adapter for dify-agent.
|
||||
|
||||
The shell-session-manager SDK owns the HTTP timeout policy for long-polling
|
||||
The built-in shellctl SDK owns the HTTP timeout policy for long-polling
|
||||
shellctl requests. This adapter stays narrowly focused on translating SDK and
|
||||
transport failures into ``ShellProviderError`` so the shell layer can return
|
||||
tool observations instead of aborting the agent loop.
|
||||
@ -16,9 +16,9 @@ import time
|
||||
from collections.abc import Awaitable
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import Protocol, TypeVar
|
||||
from typing import Protocol, TypeVar, cast
|
||||
|
||||
import httpx
|
||||
import httpx2 as httpx
|
||||
|
||||
from dify_agent.adapters.shell.protocols import (
|
||||
ShellCommandProtocol,
|
||||
@ -228,8 +228,16 @@ class ShellctlFileTransfer(ShellFileTransferProtocol):
|
||||
@dataclass(slots=True)
|
||||
class ShellctlResource(ShellResourceProtocol):
|
||||
client: ShellctlClientProtocol
|
||||
commands: ShellCommandProtocol
|
||||
files: ShellFileTransferProtocol
|
||||
_commands: ShellCommandProtocol
|
||||
_files: ShellFileTransferProtocol
|
||||
|
||||
@property
|
||||
def commands(self) -> ShellCommandProtocol:
|
||||
return self._commands
|
||||
|
||||
@property
|
||||
def files(self) -> ShellFileTransferProtocol:
|
||||
return self._files
|
||||
|
||||
async def close(self) -> None:
|
||||
try:
|
||||
@ -257,8 +265,8 @@ class ShellctlProvider(ShellProviderProtocol):
|
||||
)
|
||||
return ShellctlResource(
|
||||
client=client,
|
||||
commands=ShellctlCommands(client=client),
|
||||
files=ShellctlFileTransfer(client=client),
|
||||
_commands=ShellctlCommands(client=client),
|
||||
_files=ShellctlFileTransfer(client=client),
|
||||
)
|
||||
|
||||
|
||||
@ -269,9 +277,15 @@ def create_default_shellctl_client_factory(
|
||||
output_limit: int = _SHELLCTL_OUTPUT_LIMIT_BYTES,
|
||||
) -> ShellctlClientFactory:
|
||||
def factory() -> ShellctlClientProtocol:
|
||||
from shell_session_manager.shellctl.client import ShellctlClient
|
||||
from shellctl.client import ShellctlClient
|
||||
|
||||
return ShellctlClient(entrypoint, token=token, output_limit=output_limit)
|
||||
return cast(
|
||||
ShellctlClientProtocol,
|
||||
cast(
|
||||
object,
|
||||
ShellctlClient(entrypoint, token=token, output_limit=output_limit),
|
||||
),
|
||||
)
|
||||
|
||||
return factory
|
||||
|
||||
|
||||
@ -143,6 +143,7 @@ def download_file_from_environment(
|
||||
url=environment.url,
|
||||
auth_jwe=environment.auth_jwe,
|
||||
file=file_mapping,
|
||||
for_external=False,
|
||||
)
|
||||
if not hasattr(download_request, "filename") or not isinstance(download_request.filename, str):
|
||||
raise AgentStubTransferError("signed file download response is missing filename")
|
||||
@ -207,8 +208,10 @@ def _request_uploaded_tool_file_download_url(*, url: str, auth_jwe: str, referen
|
||||
file=AgentStubFileMapping(transfer_method="tool_file", reference=reference),
|
||||
),
|
||||
)
|
||||
if not hasattr(download_request, "download_url") or not isinstance(download_request.download_url, str):
|
||||
raise AgentStubTransferError("signed file download response is missing download_url")
|
||||
download_url = download_request.download_url
|
||||
if not isinstance(download_url, str) or not download_url:
|
||||
if not download_url:
|
||||
raise AgentStubTransferError("signed file download response is missing download_url")
|
||||
return download_url
|
||||
|
||||
|
||||
@ -133,6 +133,26 @@ def config_skills_push(
|
||||
|
||||
Pass a directory such as ./skills/researcher that contains SKILL.md. Other files in that directory are
|
||||
archived with the skill. Pushing a skill with an existing name replaces that config skill.
|
||||
|
||||
Skill directory requirements:
|
||||
|
||||
- Each PATH must be one skill directory; the directory basename is the config skill name.
|
||||
|
||||
- The directory must contain a top-level SKILL.md.
|
||||
|
||||
- SKILL.md must be non-empty UTF-8 Markdown.
|
||||
|
||||
- SKILL.md must start with YAML frontmatter matching this schema:
|
||||
|
||||
\b
|
||||
---
|
||||
name: <non-empty string>
|
||||
description: <string>
|
||||
---
|
||||
|
||||
- Symlinked files are rejected.
|
||||
|
||||
- Dependency/cache folders such as .git, __pycache__, .venv and node_modules should be manually cleared before push.
|
||||
"""
|
||||
_run_config_skills_push(paths=paths)
|
||||
|
||||
|
||||
@ -97,6 +97,7 @@ def request_agent_stub_file_download_sync(
|
||||
url: str,
|
||||
auth_jwe: str,
|
||||
file: AgentStubFileMapping,
|
||||
for_external: bool = True,
|
||||
timeout: float | httpx.Timeout = 30.0,
|
||||
sync_http_client: httpx.Client | None = None,
|
||||
):
|
||||
@ -109,12 +110,14 @@ def request_agent_stub_file_download_sync(
|
||||
url=endpoint.url,
|
||||
auth_jwe=auth_jwe,
|
||||
file=file,
|
||||
for_external=for_external,
|
||||
timeout=timeout,
|
||||
)
|
||||
return request_agent_stub_file_download_http_sync(
|
||||
base_url=endpoint.url,
|
||||
auth_jwe=auth_jwe,
|
||||
file=file,
|
||||
for_external=for_external,
|
||||
timeout=timeout,
|
||||
sync_http_client=sync_http_client,
|
||||
)
|
||||
|
||||
@ -124,6 +124,7 @@ def request_agent_stub_file_download_grpc_sync(
|
||||
url: str,
|
||||
auth_jwe: str,
|
||||
file: AgentStubFileMapping,
|
||||
for_external: bool = True,
|
||||
timeout: float | httpx.Timeout = 30.0,
|
||||
):
|
||||
"""Request one signed download URL through the gRPC Agent Stub endpoint.
|
||||
@ -144,7 +145,7 @@ def request_agent_stub_file_download_grpc_sync(
|
||||
auth_jwe=auth_jwe,
|
||||
method_name="CreateFileDownloadRequest",
|
||||
request_factory=lambda runtime: _require_conversions().proto_file_download_request(
|
||||
runtime.agent_stub_pb2, file=file
|
||||
runtime.agent_stub_pb2, file=file, for_external=for_external
|
||||
),
|
||||
response_parser=lambda response: _require_conversions().file_download_response_from_proto(response),
|
||||
timeout=timeout,
|
||||
|
||||
@ -117,13 +117,14 @@ def request_agent_stub_file_download_http_sync(
|
||||
base_url: str,
|
||||
auth_jwe: str,
|
||||
file: AgentStubFileMapping,
|
||||
for_external: bool = True,
|
||||
timeout: float | httpx.Timeout = 30.0,
|
||||
sync_http_client: httpx.Client | None = None,
|
||||
) -> AgentStubFileDownloadResponse:
|
||||
"""Request one signed download URL from the HTTP Agent Stub endpoint."""
|
||||
|
||||
try:
|
||||
request_model = AgentStubFileDownloadRequest(file=file)
|
||||
request_model = AgentStubFileDownloadRequest(file=file, for_external=for_external)
|
||||
except ValidationError as exc:
|
||||
raise AgentStubValidationError("invalid Agent Stub file download request") from exc
|
||||
response = _post_agent_stub_json(
|
||||
@ -131,7 +132,7 @@ def request_agent_stub_file_download_http_sync(
|
||||
auth_jwe=auth_jwe,
|
||||
endpoint_name="file download request",
|
||||
endpoint_url_factory=agent_stub_file_download_request_url,
|
||||
request_body=request_model.model_dump_json(exclude_none=True),
|
||||
request_body=request_model.model_dump_json(exclude_none=True, exclude_defaults=True),
|
||||
timeout=timeout,
|
||||
sync_http_client=sync_http_client,
|
||||
)
|
||||
|
||||
@ -1,6 +1,6 @@
|
||||
# pyright: reportAttributeAccessIssue=false
|
||||
# -*- coding: utf-8 -*-
|
||||
# Generated by the protocol buffer compiler. DO NOT EDIT!
|
||||
# NO CHECKED-IN PROTOBUF GENCODE
|
||||
# source: dify/agent/stub/v1/agent_stub.proto
|
||||
# Protobuf Python Version: 6.33.5
|
||||
"""Generated protocol buffer code."""
|
||||
@ -12,12 +12,7 @@ from google.protobuf import symbol_database as _symbol_database
|
||||
from google.protobuf.internal import builder as _builder
|
||||
|
||||
_runtime_version.ValidateProtobufRuntimeVersion(
|
||||
_runtime_version.Domain.PUBLIC,
|
||||
6,
|
||||
33,
|
||||
5,
|
||||
"",
|
||||
"dify/agent/stub/v1/agent_stub.proto",
|
||||
_runtime_version.Domain.PUBLIC, 6, 33, 5, "", "dify/agent/stub/v1/agent_stub.proto"
|
||||
)
|
||||
# @@protoc_insertion_point(imports)
|
||||
|
||||
@ -25,12 +20,12 @@ _sym_db = _symbol_database.Default()
|
||||
|
||||
|
||||
DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(
|
||||
b'\n#dify/agent/stub/v1/agent_stub.proto\x12\x12\x64ify.agent.stub.v1"O\n\x0e\x43onnectRequest\x12\x18\n\x10protocol_version\x18\x01 \x01(\x05\x12\x0c\n\x04\x61rgv\x18\x02 \x03(\t\x12\x15\n\rmetadata_json\x18\x03 \x01(\t"8\n\x0f\x43onnectResponse\x12\x15\n\rconnection_id\x18\x01 \x01(\t\x12\x0e\n\x06status\x18\x02 \x01(\t"7\n\x11\x46ileUploadRequest\x12\x10\n\x08\x66ilename\x18\x01 \x01(\t\x12\x10\n\x08mimetype\x18\x02 \x01(\t"(\n\x12\x46ileUploadResponse\x12\x12\n\nupload_url\x18\x01 \x01(\t"f\n\x0b\x46ileMapping\x12\x17\n\x0ftransfer_method\x18\x01 \x01(\t\x12\x16\n\treference\x18\x02 \x01(\tH\x00\x88\x01\x01\x12\x10\n\x03url\x18\x03 \x01(\tH\x01\x88\x01\x01\x42\x0c\n\n_referenceB\x06\n\x04_url"D\n\x13\x46ileDownloadRequest\x12-\n\x04\x66ile\x18\x01 \x01(\x0b\x32\x1f.dify.agent.stub.v1.FileMapping"r\n\x14\x46ileDownloadResponse\x12\x10\n\x08\x66ilename\x18\x01 \x01(\t\x12\x16\n\tmime_type\x18\x02 \x01(\tH\x00\x88\x01\x01\x12\x0c\n\x04size\x18\x03 \x01(\x03\x12\x14\n\x0c\x64ownload_url\x18\x04 \x01(\tB\x0c\n\n_mime_type2\xc0\x02\n\x10\x41gentStubService\x12R\n\x07\x43onnect\x12".dify.agent.stub.v1.ConnectRequest\x1a#.dify.agent.stub.v1.ConnectResponse\x12h\n\x17\x43reateFileUploadRequest\x12%.dify.agent.stub.v1.FileUploadRequest\x1a&.dify.agent.stub.v1.FileUploadResponse\x12n\n\x19\x43reateFileDownloadRequest\x12\'.dify.agent.stub.v1.FileDownloadRequest\x1a(.dify.agent.stub.v1.FileDownloadResponseb\x06proto3'
|
||||
b'\n#dify/agent/stub/v1/agent_stub.proto\x12\x12\x64ify.agent.stub.v1"O\n\x0e\x43onnectRequest\x12\x18\n\x10protocol_version\x18\x01 \x01(\x05\x12\x0c\n\x04\x61rgv\x18\x02 \x03(\t\x12\x15\n\rmetadata_json\x18\x03 \x01(\t"8\n\x0f\x43onnectResponse\x12\x15\n\rconnection_id\x18\x01 \x01(\t\x12\x0e\n\x06status\x18\x02 \x01(\t"7\n\x11\x46ileUploadRequest\x12\x10\n\x08\x66ilename\x18\x01 \x01(\t\x12\x10\n\x08mimetype\x18\x02 \x01(\t"(\n\x12\x46ileUploadResponse\x12\x12\n\nupload_url\x18\x01 \x01(\t"f\n\x0b\x46ileMapping\x12\x17\n\x0ftransfer_method\x18\x01 \x01(\t\x12\x16\n\treference\x18\x02 \x01(\tH\x00\x88\x01\x01\x12\x10\n\x03url\x18\x03 \x01(\tH\x01\x88\x01\x01\x42\x0c\n\n_referenceB\x06\n\x04_url"p\n\x13\x46ileDownloadRequest\x12-\n\x04\x66ile\x18\x01 \x01(\x0b\x32\x1f.dify.agent.stub.v1.FileMapping\x12\x19\n\x0c\x66or_external\x18\x02 \x01(\x08H\x00\x88\x01\x01\x42\x0f\n\r_for_external"r\n\x14\x46ileDownloadResponse\x12\x10\n\x08\x66ilename\x18\x01 \x01(\t\x12\x16\n\tmime_type\x18\x02 \x01(\tH\x00\x88\x01\x01\x12\x0c\n\x04size\x18\x03 \x01(\x03\x12\x14\n\x0c\x64ownload_url\x18\x04 \x01(\tB\x0c\n\n_mime_type2\xc0\x02\n\x10\x41gentStubService\x12R\n\x07\x43onnect\x12".dify.agent.stub.v1.ConnectRequest\x1a#.dify.agent.stub.v1.ConnectResponse\x12h\n\x17\x43reateFileUploadRequest\x12%.dify.agent.stub.v1.FileUploadRequest\x1a&.dify.agent.stub.v1.FileUploadResponse\x12n\n\x19\x43reateFileDownloadRequest\x12\'.dify.agent.stub.v1.FileDownloadRequest\x1a(.dify.agent.stub.v1.FileDownloadResponseb\x06proto3'
|
||||
)
|
||||
|
||||
_globals = globals()
|
||||
_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
|
||||
_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, "dify_agent.agent_stub.grpc._generated.agent_stub_pb2", _globals)
|
||||
_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, "dify.agent.stub.v1.agent_stub_pb2", _globals)
|
||||
if not _descriptor._USE_C_DESCRIPTORS:
|
||||
DESCRIPTOR._loaded_options = None
|
||||
_globals["_CONNECTREQUEST"]._serialized_start = 59
|
||||
@ -44,9 +39,9 @@ if not _descriptor._USE_C_DESCRIPTORS:
|
||||
_globals["_FILEMAPPING"]._serialized_start = 297
|
||||
_globals["_FILEMAPPING"]._serialized_end = 399
|
||||
_globals["_FILEDOWNLOADREQUEST"]._serialized_start = 401
|
||||
_globals["_FILEDOWNLOADREQUEST"]._serialized_end = 469
|
||||
_globals["_FILEDOWNLOADRESPONSE"]._serialized_start = 471
|
||||
_globals["_FILEDOWNLOADRESPONSE"]._serialized_end = 585
|
||||
_globals["_AGENTSTUBSERVICE"]._serialized_start = 588
|
||||
_globals["_AGENTSTUBSERVICE"]._serialized_end = 908
|
||||
_globals["_FILEDOWNLOADREQUEST"]._serialized_end = 513
|
||||
_globals["_FILEDOWNLOADRESPONSE"]._serialized_start = 515
|
||||
_globals["_FILEDOWNLOADRESPONSE"]._serialized_end = 629
|
||||
_globals["_AGENTSTUBSERVICE"]._serialized_start = 632
|
||||
_globals["_AGENTSTUBSERVICE"]._serialized_end = 952
|
||||
# @@protoc_insertion_point(module_scope)
|
||||
|
||||
@ -1,71 +1,69 @@
|
||||
from __future__ import annotations
|
||||
from google.protobuf.internal import containers as _containers
|
||||
from google.protobuf import descriptor as _descriptor
|
||||
from google.protobuf import message as _message
|
||||
from collections.abc import Iterable as _Iterable, Mapping as _Mapping
|
||||
from typing import ClassVar as _ClassVar, Optional as _Optional, Union as _Union
|
||||
|
||||
from collections.abc import Iterable
|
||||
DESCRIPTOR: _descriptor.FileDescriptor
|
||||
|
||||
from google.protobuf.message import Message
|
||||
|
||||
|
||||
class ConnectRequest(Message):
|
||||
class ConnectRequest(_message.Message):
|
||||
__slots__ = ("protocol_version", "argv", "metadata_json")
|
||||
PROTOCOL_VERSION_FIELD_NUMBER: _ClassVar[int]
|
||||
ARGV_FIELD_NUMBER: _ClassVar[int]
|
||||
METADATA_JSON_FIELD_NUMBER: _ClassVar[int]
|
||||
protocol_version: int
|
||||
argv: list[str]
|
||||
argv: _containers.RepeatedScalarFieldContainer[str]
|
||||
metadata_json: str
|
||||
def __init__(self, protocol_version: _Optional[int] = ..., argv: _Optional[_Iterable[str]] = ..., metadata_json: _Optional[str] = ...) -> None: ...
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
protocol_version: int = ...,
|
||||
argv: Iterable[str] = ...,
|
||||
metadata_json: str = ...,
|
||||
) -> None: ...
|
||||
|
||||
|
||||
class ConnectResponse(Message):
|
||||
class ConnectResponse(_message.Message):
|
||||
__slots__ = ("connection_id", "status")
|
||||
CONNECTION_ID_FIELD_NUMBER: _ClassVar[int]
|
||||
STATUS_FIELD_NUMBER: _ClassVar[int]
|
||||
connection_id: str
|
||||
status: str
|
||||
def __init__(self, connection_id: _Optional[str] = ..., status: _Optional[str] = ...) -> None: ...
|
||||
|
||||
def __init__(self, *, connection_id: str = ..., status: str = ...) -> None: ...
|
||||
|
||||
|
||||
class FileUploadRequest(Message):
|
||||
class FileUploadRequest(_message.Message):
|
||||
__slots__ = ("filename", "mimetype")
|
||||
FILENAME_FIELD_NUMBER: _ClassVar[int]
|
||||
MIMETYPE_FIELD_NUMBER: _ClassVar[int]
|
||||
filename: str
|
||||
mimetype: str
|
||||
def __init__(self, filename: _Optional[str] = ..., mimetype: _Optional[str] = ...) -> None: ...
|
||||
|
||||
def __init__(self, *, filename: str = ..., mimetype: str = ...) -> None: ...
|
||||
|
||||
|
||||
class FileUploadResponse(Message):
|
||||
class FileUploadResponse(_message.Message):
|
||||
__slots__ = ("upload_url",)
|
||||
UPLOAD_URL_FIELD_NUMBER: _ClassVar[int]
|
||||
upload_url: str
|
||||
def __init__(self, upload_url: _Optional[str] = ...) -> None: ...
|
||||
|
||||
def __init__(self, *, upload_url: str = ...) -> None: ...
|
||||
|
||||
|
||||
class FileMapping(Message):
|
||||
class FileMapping(_message.Message):
|
||||
__slots__ = ("transfer_method", "reference", "url")
|
||||
TRANSFER_METHOD_FIELD_NUMBER: _ClassVar[int]
|
||||
REFERENCE_FIELD_NUMBER: _ClassVar[int]
|
||||
URL_FIELD_NUMBER: _ClassVar[int]
|
||||
transfer_method: str
|
||||
reference: str
|
||||
url: str
|
||||
def __init__(self, transfer_method: _Optional[str] = ..., reference: _Optional[str] = ..., url: _Optional[str] = ...) -> None: ...
|
||||
|
||||
def __init__(self, *, transfer_method: str = ..., reference: str = ..., url: str = ...) -> None: ...
|
||||
def HasField(self, field_name: str) -> bool: ...
|
||||
|
||||
|
||||
class FileDownloadRequest(Message):
|
||||
class FileDownloadRequest(_message.Message):
|
||||
__slots__ = ("file", "for_external")
|
||||
FILE_FIELD_NUMBER: _ClassVar[int]
|
||||
FOR_EXTERNAL_FIELD_NUMBER: _ClassVar[int]
|
||||
file: FileMapping
|
||||
for_external: bool
|
||||
def __init__(self, file: _Optional[_Union[FileMapping, _Mapping]] = ..., for_external: _Optional[bool] = ...) -> None: ...
|
||||
|
||||
def __init__(self, *, file: FileMapping | None = ...) -> None: ...
|
||||
|
||||
|
||||
class FileDownloadResponse(Message):
|
||||
class FileDownloadResponse(_message.Message):
|
||||
__slots__ = ("filename", "mime_type", "size", "download_url")
|
||||
FILENAME_FIELD_NUMBER: _ClassVar[int]
|
||||
MIME_TYPE_FIELD_NUMBER: _ClassVar[int]
|
||||
SIZE_FIELD_NUMBER: _ClassVar[int]
|
||||
DOWNLOAD_URL_FIELD_NUMBER: _ClassVar[int]
|
||||
filename: str
|
||||
mime_type: str
|
||||
size: int
|
||||
download_url: str
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
filename: str = ...,
|
||||
mime_type: str = ...,
|
||||
size: int = ...,
|
||||
download_url: str = ...,
|
||||
) -> None: ...
|
||||
def HasField(self, field_name: str) -> bool: ...
|
||||
def __init__(self, filename: _Optional[str] = ..., mime_type: _Optional[str] = ..., size: _Optional[int] = ..., download_url: _Optional[str] = ...) -> None: ...
|
||||
|
||||
@ -98,13 +98,19 @@ def file_download_request_from_proto(message: agent_stub_pb2.FileDownloadRequest
|
||||
"reference": message.file.reference if message.file.HasField("reference") else None,
|
||||
"url": message.file.url if message.file.HasField("url") else None,
|
||||
}
|
||||
return AgentStubFileDownloadRequest.model_validate({"file": file_mapping_kwargs})
|
||||
return AgentStubFileDownloadRequest.model_validate(
|
||||
{
|
||||
"file": file_mapping_kwargs,
|
||||
"for_external": message.for_external if message.HasField("for_external") else True,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def proto_file_download_request(
|
||||
pb2_module,
|
||||
*,
|
||||
file: AgentStubFileMapping,
|
||||
for_external: bool = True,
|
||||
) -> agent_stub_pb2.FileDownloadRequest:
|
||||
"""Build one protobuf file-download request from the public DTO."""
|
||||
mapping = pb2_module.FileMapping(transfer_method=file.transfer_method)
|
||||
@ -112,7 +118,9 @@ def proto_file_download_request(
|
||||
mapping.reference = file.reference
|
||||
if file.url is not None:
|
||||
mapping.url = file.url
|
||||
return pb2_module.FileDownloadRequest(file=mapping)
|
||||
request = pb2_module.FileDownloadRequest(file=mapping)
|
||||
request.for_external = for_external
|
||||
return request
|
||||
|
||||
|
||||
def file_download_response_from_proto(message: agent_stub_pb2.FileDownloadResponse) -> AgentStubFileDownloadResponse:
|
||||
|
||||
@ -249,6 +249,7 @@ class AgentStubFileDownloadRequest(BaseModel):
|
||||
"""Request body for one signed download URL allocation."""
|
||||
|
||||
file: AgentStubFileMapping
|
||||
for_external: bool = True
|
||||
|
||||
model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid")
|
||||
|
||||
|
||||
@ -161,6 +161,8 @@ class DifyApiAgentStubFileRequestHandler:
|
||||
"invoke_from": execution_context.invoke_from,
|
||||
"file": request.file.model_dump(mode="json", exclude_none=True),
|
||||
}
|
||||
if request.for_external is False:
|
||||
payload["for_external"] = False
|
||||
data = await self._post_inner_api("/inner/api/download/file/request", payload)
|
||||
try:
|
||||
return AgentStubFileDownloadResponse.model_validate(data)
|
||||
|
||||
@ -141,7 +141,18 @@ shell_run script rules:
|
||||
|
||||
Tips:
|
||||
|
||||
- When using Python, prefer a uv script with a PEP 723 dependency header.
|
||||
- Python 3.12, uv, pip, Node.js, pnpm, and pnx are preinstalled in the local sandbox.
|
||||
- For one-off Python dependencies, prefer a uv script with a PEP 723 dependency header or:
|
||||
`uv run --with <package> python <script-or--c>`.
|
||||
- For reusable Python CLI tools, use `uv tool install <tool>`; installed commands land in `$HOME/.local/bin`.
|
||||
Run them by full path or add `$HOME/.local/bin` to PATH in the command that needs them.
|
||||
- `python3 -m pip install --user <package>` also installs into `$HOME/.local`; add `$HOME/.local/bin` to PATH
|
||||
when you need console scripts.
|
||||
- For reusable Node.js CLIs, use user-level global installs:
|
||||
`PNPM_HOME=$HOME/.local/share/pnpm PATH=$HOME/.local/share/pnpm/bin:$PATH pnpm add -g <package>`.
|
||||
Installed commands land in `$PNPM_HOME/bin`; run them by full path or with the same PATH prefix.
|
||||
- For one-off Node.js CLIs, prefer `pnx <command> [args]`.
|
||||
- Do not install new packages into system or image tool paths such as `/usr/local`, `/usr`, or `/opt/dify-agent-tools`.
|
||||
- If you need MCP, install the MCP server in the shell environment and start that server when you use it.
|
||||
|
||||
Example shell_run script:
|
||||
|
||||
@ -31,13 +31,14 @@ both the JSON-safe final output or deferred tool call and the session snapshot;
|
||||
there are no separate output or snapshot events to correlate.
|
||||
"""
|
||||
|
||||
from collections.abc import AsyncIterable, Callable
|
||||
from collections.abc import AsyncIterable, Callable, Mapping
|
||||
from collections import Counter
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Literal, Protocol, cast, runtime_checkable
|
||||
|
||||
import httpx
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
from pydantic_ai.exceptions import ModelHTTPError
|
||||
from pydantic_ai.messages import AgentStreamEvent, PartDeltaEvent, PartStartEvent, TextPart, TextPartDelta
|
||||
from pydantic_ai.output import OutputSpec
|
||||
from pydantic_ai.tools import DeferredToolRequests, DeferredToolResults
|
||||
@ -104,6 +105,28 @@ class AgentRunValidationError(ValueError):
|
||||
"""Raised when a run request is valid JSON but cannot execute."""
|
||||
|
||||
|
||||
def _run_failed_error_payload(exc: Exception) -> tuple[str, str | None]:
|
||||
"""Return the public failed-run error text and structured reason."""
|
||||
message = str(exc) or type(exc).__name__
|
||||
reason: str | None = None
|
||||
|
||||
if isinstance(exc, ModelHTTPError):
|
||||
body = exc.body
|
||||
if isinstance(body, Mapping):
|
||||
body_message = body.get("message")
|
||||
if isinstance(body_message, str) and body_message:
|
||||
message = body_message
|
||||
|
||||
error_type = body.get("error_type")
|
||||
if isinstance(error_type, str) and error_type:
|
||||
reason = error_type
|
||||
|
||||
if reason is None and exc.status_code == 429:
|
||||
reason = "InvokeRateLimitError"
|
||||
|
||||
return message, reason
|
||||
|
||||
|
||||
def _has_model_layer(request: CreateRunRequest) -> bool:
|
||||
"""Return whether the public composition includes the reserved model layer."""
|
||||
return any(layer.name == DIFY_AGENT_MODEL_LAYER_ID for layer in request.composition.layers)
|
||||
@ -165,8 +188,8 @@ class AgentRunRunner:
|
||||
try:
|
||||
outcome = await self._run_agent()
|
||||
except Exception as exc:
|
||||
message = str(exc) or type(exc).__name__
|
||||
_ = await emit_run_failed(self.sink, run_id=self.run_id, error=message)
|
||||
message, reason = _run_failed_error_payload(exc)
|
||||
_ = await emit_run_failed(self.sink, run_id=self.run_id, error=message, reason=reason)
|
||||
await self.sink.update_status(self.run_id, "failed", message)
|
||||
raise
|
||||
|
||||
|
||||
124
dify-agent/src/shellctl/__init__.py
Normal file
124
dify-agent/src/shellctl/__init__.py
Normal file
@ -0,0 +1,124 @@
|
||||
"""Public shellctl package exports.
|
||||
|
||||
This package stays lazy on purpose. Hot-path runtime helpers live outside the
|
||||
`shellctl` package, and importing this package root should
|
||||
not pull the full client/server/public DTO surface unless a caller explicitly
|
||||
asks for those exports.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from importlib import import_module
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from shellctl.client import (
|
||||
ShellctlClient,
|
||||
ShellctlClientError,
|
||||
)
|
||||
from shellctl.shared import (
|
||||
DEFAULT_AUTH_TOKEN_ENV,
|
||||
DEFAULT_BASE_URL,
|
||||
DEFAULT_BASE_URL_ENV,
|
||||
DEFAULT_GC_FINISHED_JOB_RETENTION_SECONDS,
|
||||
DEFAULT_GC_INTERVAL_SECONDS,
|
||||
DEFAULT_IDLE_FLUSH_SECONDS,
|
||||
DEFAULT_LIST_LIMIT,
|
||||
DEFAULT_OUTPUT_LIMIT_BYTES,
|
||||
DEFAULT_TERMINAL_COLS,
|
||||
DEFAULT_TERMINAL_ROWS,
|
||||
DEFAULT_TERMINATE_GRACE_SECONDS,
|
||||
DEFAULT_TIMEOUT_SECONDS,
|
||||
DeleteJobResponse,
|
||||
HealthResponse,
|
||||
InputJobRequest,
|
||||
JobInfo,
|
||||
JobResult,
|
||||
JobStatusName,
|
||||
JobStatusView,
|
||||
ListJobsResponse,
|
||||
RunJobRequest,
|
||||
TerminalSize,
|
||||
TerminateJobRequest,
|
||||
WaitJobRequest,
|
||||
generate_job_id,
|
||||
read_output_window,
|
||||
tail_output_window,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"DEFAULT_AUTH_TOKEN_ENV",
|
||||
"DEFAULT_BASE_URL",
|
||||
"DEFAULT_BASE_URL_ENV",
|
||||
"DEFAULT_GC_FINISHED_JOB_RETENTION_SECONDS",
|
||||
"DEFAULT_GC_INTERVAL_SECONDS",
|
||||
"DEFAULT_IDLE_FLUSH_SECONDS",
|
||||
"DEFAULT_LIST_LIMIT",
|
||||
"DEFAULT_OUTPUT_LIMIT_BYTES",
|
||||
"DEFAULT_TERMINAL_COLS",
|
||||
"DEFAULT_TERMINAL_ROWS",
|
||||
"DEFAULT_TERMINATE_GRACE_SECONDS",
|
||||
"DEFAULT_TIMEOUT_SECONDS",
|
||||
"DeleteJobResponse",
|
||||
"HealthResponse",
|
||||
"InputJobRequest",
|
||||
"JobInfo",
|
||||
"JobResult",
|
||||
"JobStatusName",
|
||||
"JobStatusView",
|
||||
"ListJobsResponse",
|
||||
"RunJobRequest",
|
||||
"ShellctlClient",
|
||||
"ShellctlClientError",
|
||||
"TerminalSize",
|
||||
"TerminateJobRequest",
|
||||
"WaitJobRequest",
|
||||
"generate_job_id",
|
||||
"read_output_window",
|
||||
"tail_output_window",
|
||||
]
|
||||
|
||||
_EXPORTS = {
|
||||
"ShellctlClient": "shellctl.client",
|
||||
"ShellctlClientError": "shellctl.client",
|
||||
"DEFAULT_AUTH_TOKEN_ENV": "shellctl.shared",
|
||||
"DEFAULT_BASE_URL": "shellctl.shared",
|
||||
"DEFAULT_BASE_URL_ENV": "shellctl.shared",
|
||||
"DEFAULT_GC_FINISHED_JOB_RETENTION_SECONDS": "shellctl.shared",
|
||||
"DEFAULT_GC_INTERVAL_SECONDS": "shellctl.shared",
|
||||
"DEFAULT_IDLE_FLUSH_SECONDS": "shellctl.shared",
|
||||
"DEFAULT_LIST_LIMIT": "shellctl.shared",
|
||||
"DEFAULT_OUTPUT_LIMIT_BYTES": "shellctl.shared",
|
||||
"DEFAULT_TERMINAL_COLS": "shellctl.shared",
|
||||
"DEFAULT_TERMINAL_ROWS": "shellctl.shared",
|
||||
"DEFAULT_TERMINATE_GRACE_SECONDS": "shellctl.shared",
|
||||
"DEFAULT_TIMEOUT_SECONDS": "shellctl.shared",
|
||||
"DeleteJobResponse": "shellctl.shared",
|
||||
"HealthResponse": "shellctl.shared",
|
||||
"InputJobRequest": "shellctl.shared",
|
||||
"JobInfo": "shellctl.shared",
|
||||
"JobResult": "shellctl.shared",
|
||||
"JobStatusName": "shellctl.shared",
|
||||
"JobStatusView": "shellctl.shared",
|
||||
"ListJobsResponse": "shellctl.shared",
|
||||
"RunJobRequest": "shellctl.shared",
|
||||
"TerminalSize": "shellctl.shared",
|
||||
"TerminateJobRequest": "shellctl.shared",
|
||||
"WaitJobRequest": "shellctl.shared",
|
||||
"generate_job_id": "shellctl.shared",
|
||||
"read_output_window": "shellctl.shared",
|
||||
"tail_output_window": "shellctl.shared",
|
||||
}
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
if name not in _EXPORTS:
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
module = import_module(_EXPORTS[name])
|
||||
value = getattr(module, name) # noqa: no-new-getattr lazy export proxy
|
||||
globals()[name] = value
|
||||
return value
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return sorted(set(globals()) | set(__all__))
|
||||
587
dify-agent/src/shellctl/cli.py
Normal file
587
dify-agent/src/shellctl/cli.py
Normal file
@ -0,0 +1,587 @@
|
||||
"""Typer CLI for network-backed shellctl commands.
|
||||
|
||||
Job-management commands in this module intentionally stay on the SDK side of
|
||||
the boundary: they parse CLI options, call `ShellctlClient`, and render compact
|
||||
JSON. That keeps `shellctl --help` and `shellctl run --help` free of FastAPI,
|
||||
SQLAlchemy, tmux, and local runtime bootstrap imports.
|
||||
|
||||
Only `serve` lazily imports server-side modules when that subcommand is
|
||||
actually invoked.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Awaitable, Callable
|
||||
from pathlib import Path
|
||||
from typing import NoReturn
|
||||
|
||||
import anyio
|
||||
import httpx2 as httpx
|
||||
import typer
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
from shellctl.client import ShellctlClient, ShellctlClientError
|
||||
from shellctl.shared.constants import (
|
||||
DEFAULT_AUTH_TOKEN_ENV,
|
||||
DEFAULT_BASE_URL,
|
||||
DEFAULT_BASE_URL_ENV,
|
||||
DEFAULT_GC_FINISHED_JOB_RETENTION_SECONDS,
|
||||
DEFAULT_GC_INTERVAL_SECONDS,
|
||||
DEFAULT_IDLE_FLUSH_SECONDS,
|
||||
DEFAULT_LIST_LIMIT,
|
||||
DEFAULT_OUTPUT_LIMIT_BYTES,
|
||||
DEFAULT_TERMINAL_COLS,
|
||||
DEFAULT_TERMINAL_ROWS,
|
||||
DEFAULT_TERMINATE_GRACE_SECONDS,
|
||||
DEFAULT_TIMEOUT_SECONDS,
|
||||
MAX_LIST_LIMIT,
|
||||
MAX_OUTPUT_LIMIT_BYTES,
|
||||
)
|
||||
from shellctl.shared.schemas import (
|
||||
DeleteJobResponse,
|
||||
HealthResponse,
|
||||
JobInfo,
|
||||
JobResult,
|
||||
JobStatusName,
|
||||
JobStatusView,
|
||||
RunJobRequest,
|
||||
TerminalSize,
|
||||
)
|
||||
|
||||
cli = typer.Typer(
|
||||
no_args_is_help=True,
|
||||
pretty_exceptions_enable=False,
|
||||
rich_markup_mode=None,
|
||||
)
|
||||
|
||||
|
||||
@cli.command("health")
|
||||
def health_command(
|
||||
base_url: str = typer.Option(
|
||||
DEFAULT_BASE_URL,
|
||||
"--base-url",
|
||||
envvar=DEFAULT_BASE_URL_ENV,
|
||||
help="shellctl server base URL. You can also set SHELLCTL_BASE_URL.",
|
||||
),
|
||||
auth_token: str | None = typer.Option(
|
||||
None,
|
||||
"--auth-token",
|
||||
envvar=DEFAULT_AUTH_TOKEN_ENV,
|
||||
help="Accepted for CLI consistency but ignored because /healthz is public.",
|
||||
),
|
||||
) -> None:
|
||||
"""Call the public health endpoint and report JSON."""
|
||||
|
||||
del auth_token
|
||||
|
||||
async def action(client: ShellctlClient) -> HealthResponse:
|
||||
return await client.health()
|
||||
|
||||
_run_client_action(
|
||||
base_url=base_url,
|
||||
auth_token=None,
|
||||
action=action,
|
||||
emit=_emit_model,
|
||||
)
|
||||
|
||||
|
||||
@cli.command("run")
|
||||
def run_command(
|
||||
script: str = typer.Argument(...),
|
||||
base_url: str = typer.Option(
|
||||
DEFAULT_BASE_URL,
|
||||
"--base-url",
|
||||
envvar=DEFAULT_BASE_URL_ENV,
|
||||
help="shellctl server base URL. You can also set SHELLCTL_BASE_URL.",
|
||||
),
|
||||
auth_token: str | None = typer.Option(
|
||||
None,
|
||||
"--auth-token",
|
||||
envvar=DEFAULT_AUTH_TOKEN_ENV,
|
||||
help=(
|
||||
"Bearer token value. You can also set SHELLCTL_AUTH_TOKEN. "
|
||||
"Leave it unset or empty when the server does not require auth."
|
||||
),
|
||||
),
|
||||
cwd: Path | None = typer.Option(None, "--cwd"),
|
||||
env: list[str] | None = typer.Option(None, "--env"),
|
||||
timeout: float = typer.Option(DEFAULT_TIMEOUT_SECONDS, "--timeout"),
|
||||
output_limit: int = typer.Option(DEFAULT_OUTPUT_LIMIT_BYTES, "--output-limit"),
|
||||
idle_flush_seconds: float = typer.Option(
|
||||
DEFAULT_IDLE_FLUSH_SECONDS,
|
||||
"--idle-flush-seconds",
|
||||
),
|
||||
cols: int | None = typer.Option(None, "--cols"),
|
||||
rows: int | None = typer.Option(None, "--rows"),
|
||||
) -> None:
|
||||
"""Create a job through the running shellctl server."""
|
||||
|
||||
request = _build_model(
|
||||
RunJobRequest,
|
||||
script=script,
|
||||
cwd=str(cwd) if cwd is not None else None,
|
||||
env=_parse_env(env),
|
||||
terminal=_terminal_size(cols=cols, rows=rows),
|
||||
timeout=timeout,
|
||||
output_limit=output_limit,
|
||||
idle_flush_seconds=idle_flush_seconds,
|
||||
)
|
||||
|
||||
async def action(client: ShellctlClient) -> JobResult:
|
||||
return await client.run(
|
||||
request.script,
|
||||
cwd=request.cwd,
|
||||
env=request.env,
|
||||
timeout=request.timeout,
|
||||
terminal=request.terminal,
|
||||
)
|
||||
|
||||
_run_client_action(
|
||||
base_url=base_url,
|
||||
auth_token=auth_token,
|
||||
output_limit=output_limit,
|
||||
idle_flush_seconds=idle_flush_seconds,
|
||||
action=action,
|
||||
emit=_emit_model,
|
||||
)
|
||||
|
||||
|
||||
@cli.command("wait")
|
||||
def wait_command(
|
||||
job_id: str = typer.Argument(...),
|
||||
base_url: str = typer.Option(
|
||||
DEFAULT_BASE_URL,
|
||||
"--base-url",
|
||||
envvar=DEFAULT_BASE_URL_ENV,
|
||||
help="shellctl server base URL. You can also set SHELLCTL_BASE_URL.",
|
||||
),
|
||||
auth_token: str | None = typer.Option(
|
||||
None,
|
||||
"--auth-token",
|
||||
envvar=DEFAULT_AUTH_TOKEN_ENV,
|
||||
help=(
|
||||
"Bearer token value. You can also set SHELLCTL_AUTH_TOKEN. "
|
||||
"Leave it unset or empty when the server does not require auth."
|
||||
),
|
||||
),
|
||||
offset: int = typer.Option(..., "--offset"),
|
||||
timeout: float = typer.Option(DEFAULT_TIMEOUT_SECONDS, "--timeout"),
|
||||
output_limit: int = typer.Option(DEFAULT_OUTPUT_LIMIT_BYTES, "--output-limit"),
|
||||
idle_flush_seconds: float = typer.Option(
|
||||
DEFAULT_IDLE_FLUSH_SECONDS,
|
||||
"--idle-flush-seconds",
|
||||
),
|
||||
) -> None:
|
||||
"""Wait for incremental output, completion, truncation, or timeout."""
|
||||
|
||||
async def action(client: ShellctlClient) -> JobResult:
|
||||
return await client.wait(job_id, offset=offset, timeout=timeout)
|
||||
|
||||
_run_client_action(
|
||||
base_url=base_url,
|
||||
auth_token=auth_token,
|
||||
output_limit=output_limit,
|
||||
idle_flush_seconds=idle_flush_seconds,
|
||||
action=action,
|
||||
emit=_emit_model,
|
||||
)
|
||||
|
||||
|
||||
@cli.command("status")
|
||||
def status_command(
|
||||
job_id: str = typer.Argument(...),
|
||||
base_url: str = typer.Option(
|
||||
DEFAULT_BASE_URL,
|
||||
"--base-url",
|
||||
envvar=DEFAULT_BASE_URL_ENV,
|
||||
help="shellctl server base URL. You can also set SHELLCTL_BASE_URL.",
|
||||
),
|
||||
auth_token: str | None = typer.Option(
|
||||
None,
|
||||
"--auth-token",
|
||||
envvar=DEFAULT_AUTH_TOKEN_ENV,
|
||||
help=(
|
||||
"Bearer token value. You can also set SHELLCTL_AUTH_TOKEN. "
|
||||
"Leave it unset or empty when the server does not require auth."
|
||||
),
|
||||
),
|
||||
) -> None:
|
||||
"""Materialize the current status view for one job."""
|
||||
|
||||
async def action(client: ShellctlClient) -> JobStatusView:
|
||||
return await client.status(job_id)
|
||||
|
||||
_run_client_action(
|
||||
base_url=base_url,
|
||||
auth_token=auth_token,
|
||||
action=action,
|
||||
emit=_emit_model,
|
||||
)
|
||||
|
||||
|
||||
@cli.command("list")
|
||||
def list_command(
|
||||
base_url: str = typer.Option(
|
||||
DEFAULT_BASE_URL,
|
||||
"--base-url",
|
||||
envvar=DEFAULT_BASE_URL_ENV,
|
||||
help="shellctl server base URL. You can also set SHELLCTL_BASE_URL.",
|
||||
),
|
||||
auth_token: str | None = typer.Option(
|
||||
None,
|
||||
"--auth-token",
|
||||
envvar=DEFAULT_AUTH_TOKEN_ENV,
|
||||
help=(
|
||||
"Bearer token value. You can also set SHELLCTL_AUTH_TOKEN. "
|
||||
"Leave it unset or empty when the server does not require auth."
|
||||
),
|
||||
),
|
||||
status: JobStatusName | None = typer.Option(None, "--status"),
|
||||
limit: int = typer.Option(DEFAULT_LIST_LIMIT, "--limit", min=1, max=MAX_LIST_LIMIT),
|
||||
) -> None:
|
||||
"""List recent jobs, optionally filtered by lifecycle status."""
|
||||
|
||||
async def action(client: ShellctlClient) -> list[JobInfo]:
|
||||
return await client.list_jobs(status=status, limit=limit)
|
||||
|
||||
_run_client_action(
|
||||
base_url=base_url,
|
||||
auth_token=auth_token,
|
||||
action=action,
|
||||
emit=_emit_job_list,
|
||||
)
|
||||
|
||||
|
||||
@cli.command("input")
|
||||
def input_command(
|
||||
job_id: str = typer.Argument(...),
|
||||
text: str = typer.Argument(...),
|
||||
base_url: str = typer.Option(
|
||||
DEFAULT_BASE_URL,
|
||||
"--base-url",
|
||||
envvar=DEFAULT_BASE_URL_ENV,
|
||||
help="shellctl server base URL. You can also set SHELLCTL_BASE_URL.",
|
||||
),
|
||||
auth_token: str | None = typer.Option(
|
||||
None,
|
||||
"--auth-token",
|
||||
envvar=DEFAULT_AUTH_TOKEN_ENV,
|
||||
help=(
|
||||
"Bearer token value. You can also set SHELLCTL_AUTH_TOKEN. "
|
||||
"Leave it unset or empty when the server does not require auth."
|
||||
),
|
||||
),
|
||||
offset: int = typer.Option(..., "--offset"),
|
||||
timeout: float = typer.Option(DEFAULT_TIMEOUT_SECONDS, "--timeout"),
|
||||
output_limit: int = typer.Option(DEFAULT_OUTPUT_LIMIT_BYTES, "--output-limit"),
|
||||
idle_flush_seconds: float = typer.Option(
|
||||
DEFAULT_IDLE_FLUSH_SECONDS,
|
||||
"--idle-flush-seconds",
|
||||
),
|
||||
) -> None:
|
||||
"""Send text input to a running job and wait for the next result window."""
|
||||
|
||||
async def action(client: ShellctlClient) -> JobResult:
|
||||
return await client.input(job_id, text, offset=offset, timeout=timeout)
|
||||
|
||||
_run_client_action(
|
||||
base_url=base_url,
|
||||
auth_token=auth_token,
|
||||
output_limit=output_limit,
|
||||
idle_flush_seconds=idle_flush_seconds,
|
||||
action=action,
|
||||
emit=_emit_model,
|
||||
)
|
||||
|
||||
|
||||
@cli.command("tail")
|
||||
def tail_command(
|
||||
job_id: str = typer.Argument(...),
|
||||
base_url: str = typer.Option(
|
||||
DEFAULT_BASE_URL,
|
||||
"--base-url",
|
||||
envvar=DEFAULT_BASE_URL_ENV,
|
||||
help="shellctl server base URL. You can also set SHELLCTL_BASE_URL.",
|
||||
),
|
||||
auth_token: str | None = typer.Option(
|
||||
None,
|
||||
"--auth-token",
|
||||
envvar=DEFAULT_AUTH_TOKEN_ENV,
|
||||
help=(
|
||||
"Bearer token value. You can also set SHELLCTL_AUTH_TOKEN. "
|
||||
"Leave it unset or empty when the server does not require auth."
|
||||
),
|
||||
),
|
||||
output_limit: int = typer.Option(
|
||||
DEFAULT_OUTPUT_LIMIT_BYTES,
|
||||
"--output-limit",
|
||||
min=1,
|
||||
max=MAX_OUTPUT_LIMIT_BYTES,
|
||||
),
|
||||
) -> None:
|
||||
"""Read a UTF-8-safe output tail for one job."""
|
||||
|
||||
async def action(client: ShellctlClient) -> JobResult:
|
||||
return await client.tail(job_id)
|
||||
|
||||
_run_client_action(
|
||||
base_url=base_url,
|
||||
auth_token=auth_token,
|
||||
output_limit=output_limit,
|
||||
action=action,
|
||||
emit=_emit_model,
|
||||
)
|
||||
|
||||
|
||||
@cli.command("terminate")
|
||||
def terminate_command(
|
||||
job_id: str = typer.Argument(...),
|
||||
base_url: str = typer.Option(
|
||||
DEFAULT_BASE_URL,
|
||||
"--base-url",
|
||||
envvar=DEFAULT_BASE_URL_ENV,
|
||||
help="shellctl server base URL. You can also set SHELLCTL_BASE_URL.",
|
||||
),
|
||||
auth_token: str | None = typer.Option(
|
||||
None,
|
||||
"--auth-token",
|
||||
envvar=DEFAULT_AUTH_TOKEN_ENV,
|
||||
help=(
|
||||
"Bearer token value. You can also set SHELLCTL_AUTH_TOKEN. "
|
||||
"Leave it unset or empty when the server does not require auth."
|
||||
),
|
||||
),
|
||||
grace_seconds: float = typer.Option(
|
||||
DEFAULT_TERMINATE_GRACE_SECONDS,
|
||||
"--grace-seconds",
|
||||
),
|
||||
) -> None:
|
||||
"""Terminate a job and return its materialized status."""
|
||||
|
||||
async def action(client: ShellctlClient) -> JobStatusView:
|
||||
return await client.terminate(job_id, grace_seconds=grace_seconds)
|
||||
|
||||
_run_client_action(
|
||||
base_url=base_url,
|
||||
auth_token=auth_token,
|
||||
action=action,
|
||||
emit=_emit_model,
|
||||
)
|
||||
|
||||
|
||||
@cli.command("delete")
|
||||
def delete_command(
|
||||
job_id: str = typer.Argument(...),
|
||||
base_url: str = typer.Option(
|
||||
DEFAULT_BASE_URL,
|
||||
"--base-url",
|
||||
envvar=DEFAULT_BASE_URL_ENV,
|
||||
help="shellctl server base URL. You can also set SHELLCTL_BASE_URL.",
|
||||
),
|
||||
auth_token: str | None = typer.Option(
|
||||
None,
|
||||
"--auth-token",
|
||||
envvar=DEFAULT_AUTH_TOKEN_ENV,
|
||||
help=(
|
||||
"Bearer token value. You can also set SHELLCTL_AUTH_TOKEN. "
|
||||
"Leave it unset or empty when the server does not require auth."
|
||||
),
|
||||
),
|
||||
force: bool = typer.Option(False, "--force"),
|
||||
grace_seconds: float = typer.Option(
|
||||
DEFAULT_TERMINATE_GRACE_SECONDS,
|
||||
"--grace-seconds",
|
||||
),
|
||||
) -> None:
|
||||
"""Delete a job row and artifacts, optionally terminating first."""
|
||||
|
||||
async def action(client: ShellctlClient) -> DeleteJobResponse:
|
||||
return await client.delete(
|
||||
job_id,
|
||||
force=force,
|
||||
grace_seconds=grace_seconds,
|
||||
)
|
||||
|
||||
_run_client_action(
|
||||
base_url=base_url,
|
||||
auth_token=auth_token,
|
||||
action=action,
|
||||
emit=_emit_model,
|
||||
)
|
||||
|
||||
|
||||
@cli.command("serve")
|
||||
def serve_command(
|
||||
listen: str = "127.0.0.1:8765",
|
||||
auth_token: str | None = typer.Option(
|
||||
None,
|
||||
"--auth-token",
|
||||
envvar=DEFAULT_AUTH_TOKEN_ENV,
|
||||
help=(
|
||||
"Bearer token value. You can also set SHELLCTL_AUTH_TOKEN. "
|
||||
"Leave it unset or empty to disable HTTP bearer auth."
|
||||
),
|
||||
),
|
||||
state_dir: Path | None = None,
|
||||
runtime_dir: Path | None = None,
|
||||
gc_interval_seconds: float = typer.Option(
|
||||
DEFAULT_GC_INTERVAL_SECONDS,
|
||||
"--gc-interval-seconds",
|
||||
),
|
||||
gc_finished_job_retention_seconds: float = typer.Option(
|
||||
DEFAULT_GC_FINISHED_JOB_RETENTION_SECONDS,
|
||||
"--gc-finished-job-retention-seconds",
|
||||
),
|
||||
) -> None:
|
||||
"""Run the local shellctl FastAPI server via uvicorn."""
|
||||
|
||||
from shellctl.server.serve import (
|
||||
serve_command as server_serve_command,
|
||||
)
|
||||
|
||||
server_serve_command(
|
||||
listen=listen,
|
||||
auth_token=auth_token,
|
||||
state_dir=state_dir,
|
||||
runtime_dir=runtime_dir,
|
||||
gc_interval_seconds=gc_interval_seconds,
|
||||
gc_finished_job_retention_seconds=gc_finished_job_retention_seconds,
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""CLI entrypoint used by the console script and `python -m` invocations."""
|
||||
|
||||
cli()
|
||||
|
||||
|
||||
def _parse_env(values: list[str] | None) -> dict[str, str] | None:
|
||||
if not values:
|
||||
return None
|
||||
|
||||
parsed: dict[str, str] = {}
|
||||
for value in values:
|
||||
if "=" not in value:
|
||||
raise typer.BadParameter(
|
||||
"env entries must use NAME=VALUE format",
|
||||
param_hint="--env",
|
||||
)
|
||||
name, env_value = value.split("=", 1)
|
||||
if not name:
|
||||
raise typer.BadParameter(
|
||||
"env names must be non-empty",
|
||||
param_hint="--env",
|
||||
)
|
||||
parsed[name] = env_value
|
||||
return parsed
|
||||
|
||||
|
||||
def _terminal_size(*, cols: int | None, rows: int | None) -> TerminalSize | None:
|
||||
if cols is None and rows is None:
|
||||
return None
|
||||
return _build_model(
|
||||
TerminalSize,
|
||||
cols=cols if cols is not None else DEFAULT_TERMINAL_COLS,
|
||||
rows=rows if rows is not None else DEFAULT_TERMINAL_ROWS,
|
||||
)
|
||||
|
||||
|
||||
def _build_model[ModelT: BaseModel](model_type: type[ModelT], /, **data: object) -> ModelT:
|
||||
try:
|
||||
return model_type(**data)
|
||||
except ValidationError as exc:
|
||||
raise typer.BadParameter(_validation_error_message(exc)) from exc
|
||||
|
||||
|
||||
async def _with_client[ResponseT](
|
||||
base_url: str,
|
||||
auth_token: str | None,
|
||||
output_limit: int,
|
||||
idle_flush_seconds: float,
|
||||
action: Callable[[ShellctlClient], Awaitable[ResponseT]],
|
||||
) -> ResponseT:
|
||||
async with ShellctlClient(
|
||||
base_url,
|
||||
output_limit=output_limit,
|
||||
idle_flush_seconds=idle_flush_seconds,
|
||||
token=auth_token,
|
||||
) as client:
|
||||
return await action(client)
|
||||
|
||||
|
||||
def _run_client_action[ResponseT](
|
||||
*,
|
||||
base_url: str,
|
||||
auth_token: str | None,
|
||||
action: Callable[[ShellctlClient], Awaitable[ResponseT]],
|
||||
emit: Callable[[ResponseT], None],
|
||||
output_limit: int = DEFAULT_OUTPUT_LIMIT_BYTES,
|
||||
idle_flush_seconds: float = DEFAULT_IDLE_FLUSH_SECONDS,
|
||||
) -> None:
|
||||
try:
|
||||
payload = anyio.run(
|
||||
_with_client,
|
||||
base_url,
|
||||
auth_token,
|
||||
output_limit,
|
||||
idle_flush_seconds,
|
||||
action,
|
||||
)
|
||||
except ShellctlClientError as exc:
|
||||
_emit_error_and_exit(exc.code, exc.message)
|
||||
except httpx.TimeoutException:
|
||||
_emit_error_and_exit("request_timeout", "request timed out")
|
||||
except httpx.TransportError as exc:
|
||||
_emit_error_and_exit("connection_error", str(exc))
|
||||
|
||||
emit(payload)
|
||||
|
||||
|
||||
def _emit_model(model: BaseModel) -> None:
|
||||
typer.echo(model.model_dump_json(exclude_none=True), color=False)
|
||||
|
||||
|
||||
def _emit_job_list(jobs: list[JobInfo]) -> None:
|
||||
typer.echo(
|
||||
json.dumps(
|
||||
[item.model_dump(mode="json", exclude_none=True) for item in jobs],
|
||||
separators=(",", ":"),
|
||||
),
|
||||
color=False,
|
||||
)
|
||||
|
||||
|
||||
def _emit_error_and_exit(code: str, message: str) -> NoReturn:
|
||||
typer.echo(
|
||||
json.dumps(
|
||||
{"error": {"code": code, "message": message}},
|
||||
separators=(",", ":"),
|
||||
),
|
||||
err=True,
|
||||
color=False,
|
||||
)
|
||||
raise typer.Exit(code=1)
|
||||
|
||||
|
||||
def _validation_error_message(exc: ValidationError) -> str:
|
||||
detail = exc.errors(include_url=False)[0]
|
||||
location = ".".join(str(part) for part in detail.get("loc", ()))
|
||||
message = detail["msg"]
|
||||
return f"{location}: {message}" if location else str(message)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"cli",
|
||||
"delete_command",
|
||||
"health_command",
|
||||
"input_command",
|
||||
"list_command",
|
||||
"main",
|
||||
"run_command",
|
||||
"serve_command",
|
||||
"status_command",
|
||||
"tail_command",
|
||||
"terminate_command",
|
||||
"wait_command",
|
||||
]
|
||||
12
dify-agent/src/shellctl/client/__init__.py
Normal file
12
dify-agent/src/shellctl/client/__init__.py
Normal file
@ -0,0 +1,12 @@
|
||||
"""Async HTTP client package for shellctl.
|
||||
|
||||
`shellctl.client` remains importable as before, but it is
|
||||
now a package so future client helpers can live beside the main SDK class.
|
||||
"""
|
||||
|
||||
from shellctl.client.sdk import (
|
||||
ShellctlClient,
|
||||
ShellctlClientError,
|
||||
)
|
||||
|
||||
__all__ = ["ShellctlClient", "ShellctlClientError"]
|
||||
319
dify-agent/src/shellctl/client/sdk.py
Normal file
319
dify-agent/src/shellctl/client/sdk.py
Normal file
@ -0,0 +1,319 @@
|
||||
"""Async HTTP client for the shellctl server API.
|
||||
|
||||
The SDK keeps transport-level knobs (`output_limit`, `idle_flush_seconds`, base
|
||||
URL selection, and bearer-token handling) on the client instance so individual
|
||||
method calls stay close to the network CLI's high-level workflow. Blocking shell
|
||||
operations reuse the shared client, but they override the HTTP read timeout per
|
||||
request so the transport does not fail before the server-side shell wait timeout
|
||||
or terminate-grace budget does.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
import httpx2 as httpx
|
||||
|
||||
from shellctl.shared.constants import (
|
||||
DEFAULT_AUTH_TOKEN_ENV,
|
||||
DEFAULT_BASE_URL,
|
||||
DEFAULT_IDLE_FLUSH_SECONDS,
|
||||
DEFAULT_LIST_LIMIT,
|
||||
DEFAULT_OUTPUT_LIMIT_BYTES,
|
||||
DEFAULT_TERMINATE_GRACE_SECONDS,
|
||||
DEFAULT_TIMEOUT_SECONDS,
|
||||
)
|
||||
from shellctl.shared.schemas import (
|
||||
DeleteJobResponse,
|
||||
HealthResponse,
|
||||
JobInfo,
|
||||
JobResult,
|
||||
JobStatusView,
|
||||
ListJobsResponse,
|
||||
RunJobRequest,
|
||||
TerminalSize,
|
||||
)
|
||||
|
||||
|
||||
class ShellctlClientError(RuntimeError):
|
||||
"""Raised for API-declared failures and response decode/shape problems.
|
||||
|
||||
`ShellctlClient` raises this error when the server returns a structured
|
||||
error payload, and also when an otherwise successful HTTP response contains
|
||||
invalid JSON or a top-level payload shape that does not match the SDK
|
||||
contract. Transport and timeout failures remain raw `httpx2` exceptions so
|
||||
library callers can decide how to handle network-layer failures.
|
||||
"""
|
||||
|
||||
def __init__(self, status_code: int, code: str, message: str) -> None:
|
||||
super().__init__(f"{code} ({status_code}): {message}")
|
||||
self.status_code = status_code
|
||||
self.code = code
|
||||
self.message = message
|
||||
|
||||
|
||||
class ShellctlClient:
|
||||
"""Thin async SDK for the shellctl HTTP API.
|
||||
|
||||
The client owns a reusable `httpx.AsyncClient` unless one is injected via the
|
||||
`client` argument. Callers can therefore either keep one instance for a full
|
||||
workflow or treat it as an async context manager. Injected clients keep their
|
||||
original lifecycle; `close()` only closes clients that this SDK created.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
base_url: str = DEFAULT_BASE_URL,
|
||||
*,
|
||||
output_limit: int = DEFAULT_OUTPUT_LIMIT_BYTES,
|
||||
idle_flush_seconds: float = DEFAULT_IDLE_FLUSH_SECONDS,
|
||||
token: str | None = None,
|
||||
client: httpx.AsyncClient | None = None,
|
||||
transport: httpx.AsyncBaseTransport | None = None,
|
||||
request_timeout_grace_seconds: float = 10.0,
|
||||
) -> None:
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.output_limit = output_limit
|
||||
self.idle_flush_seconds = idle_flush_seconds
|
||||
self.request_timeout_grace_seconds = request_timeout_grace_seconds
|
||||
self.token = token if token is not None else os.environ.get(DEFAULT_AUTH_TOKEN_ENV)
|
||||
self._owns_client = client is None
|
||||
self._client = client or httpx.AsyncClient(
|
||||
base_url=self.base_url,
|
||||
follow_redirects=True,
|
||||
timeout=httpx.Timeout(DEFAULT_TIMEOUT_SECONDS, connect=DEFAULT_TIMEOUT_SECONDS),
|
||||
transport=transport,
|
||||
)
|
||||
|
||||
async def __aenter__(self) -> ShellctlClient:
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None:
|
||||
await self.close()
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Close the underlying HTTP client if this SDK instance owns it."""
|
||||
|
||||
if self._owns_client:
|
||||
await self._client.aclose()
|
||||
|
||||
def _wait_request_timeout(self, timeout: float) -> httpx.Timeout:
|
||||
"""Return a request timeout for blocking shell calls.
|
||||
|
||||
The shellctl server enforces the payload timeout, while the SDK keeps a
|
||||
small HTTP read-timeout grace so the transport can wait slightly longer
|
||||
for that response without loosening connect/write/pool timeouts.
|
||||
"""
|
||||
|
||||
return httpx.Timeout(
|
||||
connect=DEFAULT_TIMEOUT_SECONDS,
|
||||
read=timeout + self.request_timeout_grace_seconds,
|
||||
write=DEFAULT_TIMEOUT_SECONDS,
|
||||
pool=DEFAULT_TIMEOUT_SECONDS,
|
||||
)
|
||||
|
||||
def _terminate_request_timeout(self, grace_seconds: float | None = None) -> httpx.Timeout:
|
||||
"""Return a request timeout for terminate-style calls.
|
||||
|
||||
`terminate()` and forced `delete()` block until the server finishes the
|
||||
terminate grace window, so their HTTP read timeout must cover that
|
||||
business wait budget even when the request relies on the API default.
|
||||
"""
|
||||
|
||||
effective_grace_seconds = DEFAULT_TERMINATE_GRACE_SECONDS if grace_seconds is None else grace_seconds
|
||||
return self._wait_request_timeout(effective_grace_seconds)
|
||||
|
||||
async def health(self) -> HealthResponse:
|
||||
"""Call the public health endpoint and decode it as `HealthResponse`."""
|
||||
|
||||
return HealthResponse.model_validate(await self.healthz())
|
||||
|
||||
async def healthz(self) -> dict[str, Any]:
|
||||
"""Call the public health endpoint without requiring auth."""
|
||||
|
||||
response = await self._client.get("/healthz")
|
||||
return self._decode_response(response)
|
||||
|
||||
async def run(
|
||||
self,
|
||||
script: str,
|
||||
*,
|
||||
cwd: str | None = None,
|
||||
env: dict[str, str] | None = None,
|
||||
timeout: float = DEFAULT_TIMEOUT_SECONDS,
|
||||
terminal: TerminalSize | None = None,
|
||||
) -> JobResult:
|
||||
"""Create a new job and wait for initial output or completion.
|
||||
|
||||
`cwd` and `env` preset the script's working directory and environment
|
||||
overlay on the server side.
|
||||
"""
|
||||
|
||||
payload = RunJobRequest(
|
||||
script=script,
|
||||
cwd=cwd,
|
||||
env=env,
|
||||
terminal=terminal,
|
||||
timeout=timeout,
|
||||
output_limit=self.output_limit,
|
||||
idle_flush_seconds=self.idle_flush_seconds,
|
||||
)
|
||||
response = await self._client.post(
|
||||
"/v1/jobs/run",
|
||||
json=payload.model_dump(mode="json", exclude_none=True),
|
||||
headers=self._auth_headers(),
|
||||
timeout=self._wait_request_timeout(timeout),
|
||||
)
|
||||
return JobResult.model_validate(self._decode_response(response))
|
||||
|
||||
async def wait(
|
||||
self,
|
||||
job_id: str,
|
||||
*,
|
||||
offset: int,
|
||||
timeout: float = DEFAULT_TIMEOUT_SECONDS,
|
||||
) -> JobResult:
|
||||
"""Wait for incremental output, completion, truncation, or timeout."""
|
||||
|
||||
response = await self._client.post(
|
||||
f"/v1/jobs/{job_id}/wait",
|
||||
json={
|
||||
"offset": offset,
|
||||
"timeout": timeout,
|
||||
"output_limit": self.output_limit,
|
||||
"idle_flush_seconds": self.idle_flush_seconds,
|
||||
},
|
||||
headers=self._auth_headers(),
|
||||
timeout=self._wait_request_timeout(timeout),
|
||||
)
|
||||
return JobResult.model_validate(self._decode_response(response))
|
||||
|
||||
async def status(self, job_id: str) -> JobStatusView:
|
||||
"""Fetch the materialized status view for one job."""
|
||||
|
||||
response = await self._client.get(
|
||||
f"/v1/jobs/{job_id}",
|
||||
headers=self._auth_headers(),
|
||||
)
|
||||
return JobStatusView.model_validate(self._decode_response(response))
|
||||
|
||||
async def list_jobs(
|
||||
self,
|
||||
*,
|
||||
status: str | None = None,
|
||||
limit: int = DEFAULT_LIST_LIMIT,
|
||||
) -> list[JobInfo]:
|
||||
"""List recent jobs, optionally filtered by lifecycle status."""
|
||||
|
||||
params: dict[str, Any] = {"limit": limit}
|
||||
if status is not None:
|
||||
params["status"] = status
|
||||
response = await self._client.get(
|
||||
"/v1/jobs",
|
||||
params=params,
|
||||
headers=self._auth_headers(),
|
||||
)
|
||||
payload = ListJobsResponse.model_validate(self._decode_response(response))
|
||||
return payload.jobs
|
||||
|
||||
async def input(
|
||||
self,
|
||||
job_id: str,
|
||||
text: str,
|
||||
*,
|
||||
offset: int,
|
||||
timeout: float = DEFAULT_TIMEOUT_SECONDS,
|
||||
) -> JobResult:
|
||||
"""Send text input to a running job and then wait like `wait()`."""
|
||||
|
||||
response = await self._client.post(
|
||||
f"/v1/jobs/{job_id}/input",
|
||||
json={
|
||||
"text": text,
|
||||
"offset": offset,
|
||||
"timeout": timeout,
|
||||
"output_limit": self.output_limit,
|
||||
"idle_flush_seconds": self.idle_flush_seconds,
|
||||
},
|
||||
headers=self._auth_headers(),
|
||||
timeout=self._wait_request_timeout(timeout),
|
||||
)
|
||||
return JobResult.model_validate(self._decode_response(response))
|
||||
|
||||
async def tail(self, job_id: str) -> JobResult:
|
||||
"""Fetch an immediate UTF-8-safe tail snapshot for a job."""
|
||||
|
||||
response = await self._client.get(
|
||||
f"/v1/jobs/{job_id}/log/tail",
|
||||
params={"output_limit": self.output_limit},
|
||||
headers=self._auth_headers(),
|
||||
)
|
||||
return JobResult.model_validate(self._decode_response(response))
|
||||
|
||||
async def terminate(
|
||||
self,
|
||||
job_id: str,
|
||||
grace_seconds: float = DEFAULT_TERMINATE_GRACE_SECONDS,
|
||||
) -> JobStatusView:
|
||||
"""Terminate a job, waiting long enough for the grace window to finish."""
|
||||
|
||||
response = await self._client.post(
|
||||
f"/v1/jobs/{job_id}/terminate",
|
||||
json={"grace_seconds": grace_seconds},
|
||||
headers=self._auth_headers(),
|
||||
timeout=self._terminate_request_timeout(grace_seconds),
|
||||
)
|
||||
return JobStatusView.model_validate(self._decode_response(response))
|
||||
|
||||
async def delete(
|
||||
self,
|
||||
job_id: str,
|
||||
*,
|
||||
force: bool = False,
|
||||
grace_seconds: float | None = None,
|
||||
) -> DeleteJobResponse:
|
||||
"""Delete job artifacts, optionally waiting for forced termination first."""
|
||||
|
||||
params: dict[str, Any] = {"force": str(force).lower()}
|
||||
if grace_seconds is not None:
|
||||
params["grace_seconds"] = grace_seconds
|
||||
request_kwargs: dict[str, Any] = {
|
||||
"params": params,
|
||||
"headers": self._auth_headers(),
|
||||
}
|
||||
if force:
|
||||
request_kwargs["timeout"] = self._terminate_request_timeout(grace_seconds)
|
||||
response = await self._client.delete(
|
||||
f"/v1/jobs/{job_id}",
|
||||
**request_kwargs,
|
||||
)
|
||||
return DeleteJobResponse.model_validate(self._decode_response(response))
|
||||
|
||||
def _auth_headers(self) -> dict[str, str]:
|
||||
if not self.token:
|
||||
return {}
|
||||
return {"Authorization": f"Bearer {self.token}"}
|
||||
|
||||
def _decode_response(self, response: httpx.Response) -> dict[str, Any]:
|
||||
try:
|
||||
payload = response.json()
|
||||
except ValueError as exc: # pragma: no cover - network/proxy corruption
|
||||
raise ShellctlClientError(response.status_code, "invalid_json", response.text) from exc
|
||||
|
||||
if response.is_error:
|
||||
error = payload.get("error") if isinstance(payload, dict) else None
|
||||
if isinstance(error, dict):
|
||||
code = str(error.get("code", "request_failed"))
|
||||
message = str(error.get("message", response.text))
|
||||
else:
|
||||
code = "request_failed"
|
||||
message = response.text
|
||||
raise ShellctlClientError(response.status_code, code, message)
|
||||
|
||||
if not isinstance(payload, dict):
|
||||
raise ShellctlClientError(response.status_code, "invalid_payload", response.text)
|
||||
return payload
|
||||
|
||||
|
||||
__all__ = ["ShellctlClient", "ShellctlClientError"]
|
||||
1
dify-agent/src/shellctl/py.typed
Normal file
1
dify-agent/src/shellctl/py.typed
Normal file
@ -0,0 +1 @@
|
||||
|
||||
58
dify-agent/src/shellctl/server/__init__.py
Normal file
58
dify-agent/src/shellctl/server/__init__.py
Normal file
@ -0,0 +1,58 @@
|
||||
"""shellctl server package.
|
||||
|
||||
The server stack sits behind lazy exports so importing the network CLI does not
|
||||
pull in FastAPI, SQLAlchemy, tmux, or the local service runtime unless a
|
||||
server-side symbol is actually used.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from importlib import import_module
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from shellctl.server.api import create_app
|
||||
from shellctl.server.cli import (
|
||||
cli,
|
||||
main,
|
||||
)
|
||||
from shellctl.server.config import ShellctlConfig
|
||||
from shellctl.server.db import JobRow
|
||||
from shellctl.server.errors import ShellctlServerError
|
||||
from shellctl.server.serve import serve_command
|
||||
from shellctl.server.service import ShellctlService
|
||||
|
||||
__all__ = [
|
||||
"JobRow",
|
||||
"ShellctlConfig",
|
||||
"ShellctlServerError",
|
||||
"ShellctlService",
|
||||
"cli",
|
||||
"create_app",
|
||||
"main",
|
||||
"serve_command",
|
||||
]
|
||||
|
||||
_EXPORTS = {
|
||||
"JobRow": "shellctl.server.db",
|
||||
"ShellctlConfig": "shellctl.server.config",
|
||||
"ShellctlServerError": "shellctl.server.errors",
|
||||
"ShellctlService": "shellctl.server.service",
|
||||
"cli": "shellctl.server.cli",
|
||||
"create_app": "shellctl.server.api",
|
||||
"main": "shellctl.server.cli",
|
||||
"serve_command": "shellctl.server.serve",
|
||||
}
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
if name not in _EXPORTS:
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
module = import_module(_EXPORTS[name])
|
||||
value = getattr(module, name) # noqa: no-new-getattr lazy export proxy
|
||||
globals()[name] = value
|
||||
return value
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return sorted(set(globals()) | set(__all__))
|
||||
6
dify-agent/src/shellctl/server/__main__.py
Normal file
6
dify-agent/src/shellctl/server/__main__.py
Normal file
@ -0,0 +1,6 @@
|
||||
"""Support `python -m shellctl.server`."""
|
||||
|
||||
from shellctl.cli import main
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
198
dify-agent/src/shellctl/server/api.py
Normal file
198
dify-agent/src/shellctl/server/api.py
Normal file
@ -0,0 +1,198 @@
|
||||
"""FastAPI wiring for shellctl server endpoints."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Annotated, cast
|
||||
|
||||
from fastapi import Depends, FastAPI, Header, Query, Request
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
from shellctl.server.config import ShellctlConfig
|
||||
from shellctl.server.errors import ShellctlServerError
|
||||
from shellctl.server.service import ShellctlService
|
||||
from shellctl.shared.constants import (
|
||||
DEFAULT_HEALTH_STATUS,
|
||||
DEFAULT_LIST_LIMIT,
|
||||
DEFAULT_OUTPUT_LIMIT_BYTES,
|
||||
DEFAULT_TERMINATE_GRACE_SECONDS,
|
||||
MAX_LIST_LIMIT,
|
||||
MAX_OUTPUT_LIMIT_BYTES,
|
||||
)
|
||||
from shellctl.shared.schemas import (
|
||||
DeleteJobResponse,
|
||||
ErrorDetail,
|
||||
ErrorResponse,
|
||||
HealthResponse,
|
||||
InputJobRequest,
|
||||
JobResult,
|
||||
JobStatusName,
|
||||
JobStatusView,
|
||||
ListJobsResponse,
|
||||
RunJobRequest,
|
||||
TerminateJobRequest,
|
||||
WaitJobRequest,
|
||||
)
|
||||
|
||||
|
||||
def create_app(
|
||||
config: ShellctlConfig | None = None,
|
||||
*,
|
||||
service: ShellctlService | None = None,
|
||||
) -> FastAPI:
|
||||
"""Create the FastAPI application used by `shellctl serve`."""
|
||||
|
||||
resolved_config = config or ShellctlConfig()
|
||||
resolved_service = service or ShellctlService(resolved_config)
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(_app: FastAPI):
|
||||
await resolved_service.initialize()
|
||||
resolved_service.start_background_gc()
|
||||
resolved_service.start_background_pipe_monitor()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
await resolved_service.shutdown()
|
||||
|
||||
app = FastAPI(title="shellctl", version="0.1.0", lifespan=lifespan)
|
||||
app.state.shellctl_service = resolved_service
|
||||
|
||||
@app.exception_handler(ShellctlServerError)
|
||||
async def handle_shellctl_error(
|
||||
_request: Request,
|
||||
exc: ShellctlServerError,
|
||||
) -> JSONResponse:
|
||||
return JSONResponse(
|
||||
status_code=exc.status_code,
|
||||
content=ErrorResponse(error=ErrorDetail(code=exc.code, message=exc.message)).model_dump(mode="json"),
|
||||
)
|
||||
|
||||
@app.exception_handler(RuntimeError)
|
||||
async def handle_runtime_error(_request: Request, exc: RuntimeError) -> JSONResponse:
|
||||
return JSONResponse(
|
||||
status_code=500,
|
||||
content=ErrorResponse(
|
||||
error=ErrorDetail(
|
||||
code="internal_error",
|
||||
message=str(exc) or "internal server error",
|
||||
)
|
||||
).model_dump(mode="json"),
|
||||
)
|
||||
|
||||
def get_service() -> ShellctlService:
|
||||
return cast(ShellctlService, app.state.shellctl_service)
|
||||
|
||||
def verify_auth(
|
||||
authorization: Annotated[str | None, Header()] = None,
|
||||
) -> None:
|
||||
token = resolved_config.auth_token
|
||||
if token is None:
|
||||
return
|
||||
expected = f"Bearer {token}"
|
||||
if authorization != expected:
|
||||
raise ShellctlServerError(401, "unauthorized", "Missing or invalid bearer token")
|
||||
|
||||
@app.get("/healthz", response_model=HealthResponse)
|
||||
async def healthz() -> HealthResponse:
|
||||
return HealthResponse(status=DEFAULT_HEALTH_STATUS)
|
||||
|
||||
@app.post(
|
||||
"/v1/jobs/run",
|
||||
response_model=JobResult,
|
||||
dependencies=[Depends(verify_auth)],
|
||||
)
|
||||
async def run_job(
|
||||
payload: RunJobRequest,
|
||||
svc: ShellctlService = Depends(get_service),
|
||||
) -> JobResult:
|
||||
return await svc.run_job(payload)
|
||||
|
||||
@app.post(
|
||||
"/v1/jobs/{job_id}/wait",
|
||||
response_model=JobResult,
|
||||
dependencies=[Depends(verify_auth)],
|
||||
)
|
||||
async def wait_job(
|
||||
job_id: str,
|
||||
payload: WaitJobRequest,
|
||||
svc: ShellctlService = Depends(get_service),
|
||||
) -> JobResult:
|
||||
return await svc.wait_job(job_id, payload)
|
||||
|
||||
@app.get(
|
||||
"/v1/jobs/{job_id}/log/tail",
|
||||
response_model=JobResult,
|
||||
dependencies=[Depends(verify_auth)],
|
||||
)
|
||||
async def tail_job(
|
||||
job_id: str,
|
||||
output_limit: Annotated[int, Query(ge=1, le=MAX_OUTPUT_LIMIT_BYTES)] = DEFAULT_OUTPUT_LIMIT_BYTES,
|
||||
svc: ShellctlService = Depends(get_service),
|
||||
) -> JobResult:
|
||||
return await svc.tail_job(job_id, output_limit=output_limit)
|
||||
|
||||
@app.get(
|
||||
"/v1/jobs/{job_id}",
|
||||
response_model=JobStatusView,
|
||||
dependencies=[Depends(verify_auth)],
|
||||
)
|
||||
async def job_status(
|
||||
job_id: str,
|
||||
svc: ShellctlService = Depends(get_service),
|
||||
) -> JobStatusView:
|
||||
return await svc.get_job_status(job_id)
|
||||
|
||||
@app.get(
|
||||
"/v1/jobs",
|
||||
response_model=ListJobsResponse,
|
||||
dependencies=[Depends(verify_auth)],
|
||||
)
|
||||
async def list_jobs(
|
||||
status: Annotated[JobStatusName | None, Query()] = None,
|
||||
limit: Annotated[int, Query(ge=1, le=MAX_LIST_LIMIT)] = DEFAULT_LIST_LIMIT,
|
||||
svc: ShellctlService = Depends(get_service),
|
||||
) -> ListJobsResponse:
|
||||
return await svc.list_jobs(status=status, limit=limit)
|
||||
|
||||
@app.post(
|
||||
"/v1/jobs/{job_id}/input",
|
||||
response_model=JobResult,
|
||||
dependencies=[Depends(verify_auth)],
|
||||
)
|
||||
async def input_job(
|
||||
job_id: str,
|
||||
payload: InputJobRequest,
|
||||
svc: ShellctlService = Depends(get_service),
|
||||
) -> JobResult:
|
||||
return await svc.send_input(job_id, payload)
|
||||
|
||||
@app.post(
|
||||
"/v1/jobs/{job_id}/terminate",
|
||||
response_model=JobStatusView,
|
||||
dependencies=[Depends(verify_auth)],
|
||||
)
|
||||
async def terminate_job(
|
||||
job_id: str,
|
||||
payload: TerminateJobRequest,
|
||||
svc: ShellctlService = Depends(get_service),
|
||||
) -> JobStatusView:
|
||||
return await svc.terminate_job(job_id, payload)
|
||||
|
||||
@app.delete(
|
||||
"/v1/jobs/{job_id}",
|
||||
response_model=DeleteJobResponse,
|
||||
dependencies=[Depends(verify_auth)],
|
||||
)
|
||||
async def delete_job(
|
||||
job_id: str,
|
||||
force: bool = False,
|
||||
grace_seconds: float = DEFAULT_TERMINATE_GRACE_SECONDS,
|
||||
svc: ShellctlService = Depends(get_service),
|
||||
) -> DeleteJobResponse:
|
||||
return await svc.delete_job(job_id, force=force, grace_seconds=grace_seconds)
|
||||
|
||||
return app
|
||||
|
||||
|
||||
__all__ = ["create_app"]
|
||||
62
dify-agent/src/shellctl/server/artifacts.py
Normal file
62
dify-agent/src/shellctl/server/artifacts.py
Normal file
@ -0,0 +1,62 @@
|
||||
"""Per-job artifact names used by shellctl server/runtime code.
|
||||
|
||||
Normal job completion is coordinated through small marker files inside each
|
||||
`jobs/<job_id>/` directory so the tmux output-pipe finalizer can publish the
|
||||
SQLite `exited(exit_code, ended_at)` state only after PTY output is fully
|
||||
drained into `output.log`. The same artifact directory also stores the request's
|
||||
environment overlay so the runner can merge arbitrary key/value pairs without
|
||||
shell-escaping them into the generated script. Separate failure markers and a
|
||||
dedicated `pipe-error.log` stderr capture keep startup diagnostics available
|
||||
when the sanitizer never reaches its ready-file handshake.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
RUNNER_EXIT_CODE_FILENAME = ".runner-exit-code"
|
||||
RUNNER_ENDED_AT_FILENAME = ".runner-ended-at"
|
||||
JOB_ENV_FILENAME = ".job-env.json"
|
||||
PIPE_DRAINED_FILENAME = ".pipe-drained"
|
||||
PIPE_FAILED_FILENAME = ".pipe-failed"
|
||||
PIPE_ERROR_LOG_FILENAME = "pipe-error.log"
|
||||
|
||||
|
||||
def runner_exit_code_path(job_dir: Path) -> Path:
|
||||
return job_dir / RUNNER_EXIT_CODE_FILENAME
|
||||
|
||||
|
||||
def runner_ended_at_path(job_dir: Path) -> Path:
|
||||
return job_dir / RUNNER_ENDED_AT_FILENAME
|
||||
|
||||
|
||||
def job_env_path(job_dir: Path) -> Path:
|
||||
return job_dir / JOB_ENV_FILENAME
|
||||
|
||||
|
||||
def pipe_drained_path(job_dir: Path) -> Path:
|
||||
return job_dir / PIPE_DRAINED_FILENAME
|
||||
|
||||
|
||||
def pipe_failed_path(job_dir: Path) -> Path:
|
||||
return job_dir / PIPE_FAILED_FILENAME
|
||||
|
||||
|
||||
def pipe_error_log_path(job_dir: Path) -> Path:
|
||||
return job_dir / PIPE_ERROR_LOG_FILENAME
|
||||
|
||||
|
||||
__all__ = [
|
||||
"JOB_ENV_FILENAME",
|
||||
"PIPE_DRAINED_FILENAME",
|
||||
"PIPE_ERROR_LOG_FILENAME",
|
||||
"PIPE_FAILED_FILENAME",
|
||||
"RUNNER_ENDED_AT_FILENAME",
|
||||
"RUNNER_EXIT_CODE_FILENAME",
|
||||
"job_env_path",
|
||||
"pipe_drained_path",
|
||||
"pipe_error_log_path",
|
||||
"pipe_failed_path",
|
||||
"runner_ended_at_path",
|
||||
"runner_exit_code_path",
|
||||
]
|
||||
17
dify-agent/src/shellctl/server/cli.py
Normal file
17
dify-agent/src/shellctl/server/cli.py
Normal file
@ -0,0 +1,17 @@
|
||||
"""Compatibility shim for historical `shellctl.server.cli` imports.
|
||||
|
||||
Job-management commands now live in `shellctl.cli`. This
|
||||
module is intentionally a thin legacy re-export for callers that still import
|
||||
CLI symbols from `shellctl.server.cli`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from shellctl.cli import cli, main
|
||||
from shellctl.server.serve import serve_command
|
||||
|
||||
__all__ = [
|
||||
"cli",
|
||||
"main",
|
||||
"serve_command",
|
||||
]
|
||||
105
dify-agent/src/shellctl/server/config.py
Normal file
105
dify-agent/src/shellctl/server/config.py
Normal file
@ -0,0 +1,105 @@
|
||||
"""Configuration objects for shellctl server/runtime modules."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import cast
|
||||
|
||||
from shellctl.shared.constants import (
|
||||
DEFAULT_AUTH_TOKEN_ENV,
|
||||
DEFAULT_GC_FINISHED_JOB_RETENTION_SECONDS,
|
||||
DEFAULT_GC_INTERVAL_SECONDS,
|
||||
DEFAULT_IDLE_FLUSH_SECONDS,
|
||||
DEFAULT_LIST_LIMIT,
|
||||
DEFAULT_OUTPUT_LIMIT_BYTES,
|
||||
DEFAULT_TERMINAL_COLS,
|
||||
DEFAULT_TERMINAL_ROWS,
|
||||
DEFAULT_TERMINATE_GRACE_SECONDS,
|
||||
DEFAULT_TIMEOUT_SECONDS,
|
||||
MAX_LIST_LIMIT,
|
||||
MAX_OUTPUT_LIMIT_BYTES,
|
||||
MAX_WAIT_TIMEOUT_SECONDS,
|
||||
)
|
||||
from shellctl.shared.runtime import (
|
||||
default_runtime_dir,
|
||||
default_state_dir,
|
||||
)
|
||||
from shellctl_runtime.paths import DEFAULT_SQLITE_BUSY_TIMEOUT_MS
|
||||
|
||||
|
||||
@dataclass(slots=True, frozen=True)
|
||||
class ShellctlConfig:
|
||||
"""Runtime configuration for the shellctl service and CLI.
|
||||
|
||||
The tmux subprocess hooks use dedicated console scripts instead of
|
||||
`python -m shellctl...` entrypoints so every job does not pay
|
||||
the shellctl client/server import cost just to sanitize PTY bytes or record
|
||||
an exit row. `sqlite_busy_timeout_ms` still applies to the out-of-process
|
||||
`runner-exit` callback, so non-default deployments keep one SQLite timeout
|
||||
policy across the service and tmux finalizer.
|
||||
|
||||
Bearer auth is opt-in: if the explicit `auth_token` and the fallback
|
||||
`SHELLCTL_AUTH_TOKEN` environment variable are both missing or empty,
|
||||
`shellctl serve` accepts requests without checking an Authorization header.
|
||||
"""
|
||||
|
||||
listen: str = "127.0.0.1:8765"
|
||||
auth_token: str | None = None
|
||||
state_dir: Path = field(default_factory=default_state_dir)
|
||||
runtime_dir: Path | None = None
|
||||
gc_interval_seconds: float = DEFAULT_GC_INTERVAL_SECONDS
|
||||
gc_finished_job_retention_seconds: float = DEFAULT_GC_FINISHED_JOB_RETENTION_SECONDS
|
||||
default_timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS
|
||||
max_wait_timeout_seconds: float = MAX_WAIT_TIMEOUT_SECONDS
|
||||
idle_flush_seconds: float = DEFAULT_IDLE_FLUSH_SECONDS
|
||||
default_cwd: Path = field(default_factory=Path.home)
|
||||
default_terminal_cols: int = DEFAULT_TERMINAL_COLS
|
||||
default_terminal_rows: int = DEFAULT_TERMINAL_ROWS
|
||||
default_list_limit: int = DEFAULT_LIST_LIMIT
|
||||
max_list_limit: int = MAX_LIST_LIMIT
|
||||
default_output_limit_bytes: int = DEFAULT_OUTPUT_LIMIT_BYTES
|
||||
max_output_limit_bytes: int = MAX_OUTPUT_LIMIT_BYTES
|
||||
default_terminate_grace_seconds: float = DEFAULT_TERMINATE_GRACE_SECONDS
|
||||
poll_interval_seconds: float = 0.05
|
||||
pipe_monitor_interval_seconds: float = 1.0
|
||||
pipe_ready_timeout_seconds: float = 10.0
|
||||
sqlite_busy_timeout_ms: int = DEFAULT_SQLITE_BUSY_TIMEOUT_MS
|
||||
sanitize_pty_command: tuple[str, ...] = ("shellctl-sanitize-pty",)
|
||||
runner_exit_command: tuple[str, ...] = ("shellctl-runner-exit",)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.runtime_dir is None:
|
||||
object.__setattr__(self, "runtime_dir", default_runtime_dir(self.state_dir))
|
||||
token = self.auth_token
|
||||
if token is None:
|
||||
token = os.environ.get(DEFAULT_AUTH_TOKEN_ENV)
|
||||
if not token:
|
||||
token = None
|
||||
object.__setattr__(self, "auth_token", token)
|
||||
|
||||
@property
|
||||
def jobs_dir(self) -> Path:
|
||||
return self.state_dir / "jobs"
|
||||
|
||||
@property
|
||||
def db_path(self) -> Path:
|
||||
return self.state_dir / "shellctl.db"
|
||||
|
||||
@property
|
||||
def database_url(self) -> str:
|
||||
return f"sqlite+aiosqlite:///{self.db_path}"
|
||||
|
||||
@property
|
||||
def tmux_socket(self) -> Path:
|
||||
runtime_dir = cast(Path, self.runtime_dir)
|
||||
return runtime_dir / "tmux.sock"
|
||||
|
||||
@property
|
||||
def runner_path(self) -> Path:
|
||||
runtime_dir = cast(Path, self.runtime_dir)
|
||||
return runtime_dir / "bin" / "shellctl-runner"
|
||||
|
||||
|
||||
__all__ = ["ShellctlConfig"]
|
||||
52
dify-agent/src/shellctl/server/db.py
Normal file
52
dify-agent/src/shellctl/server/db.py
Normal file
@ -0,0 +1,52 @@
|
||||
"""SQLite models and engine helpers for shellctl."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, cast
|
||||
|
||||
from sqlalchemy import event
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine
|
||||
from sqlmodel import Field, SQLModel
|
||||
|
||||
|
||||
class JobRow(SQLModel, table=True):
|
||||
"""SQLite source-of-truth row for shellctl jobs.
|
||||
|
||||
The table intentionally combines immutable metadata, mutable lifecycle state,
|
||||
and exit facts so state transitions can be expressed as single conditional
|
||||
`UPDATE` statements without synchronizing separate records.
|
||||
"""
|
||||
|
||||
__tablename__ = cast(Any, "jobs")
|
||||
|
||||
job_id: str = Field(primary_key=True)
|
||||
script_path: str
|
||||
output_path: str
|
||||
cwd: str
|
||||
terminal_cols: int
|
||||
terminal_rows: int
|
||||
status: str = Field(index=True)
|
||||
session_name: str
|
||||
pane_target: str
|
||||
exit_code: int | None = Field(default=None, nullable=True)
|
||||
reason: str | None = Field(default=None, nullable=True)
|
||||
message: str | None = Field(default=None, nullable=True)
|
||||
created_at: str = Field(index=True)
|
||||
started_at: str | None = Field(default=None, nullable=True)
|
||||
ended_at: str | None = Field(default=None, nullable=True, index=True)
|
||||
updated_at: str
|
||||
|
||||
|
||||
def configure_sqlite_engine(engine: AsyncEngine, *, busy_timeout_ms: int) -> None:
|
||||
"""Install SQLite pragmas required by the proposal's concurrency model."""
|
||||
|
||||
@event.listens_for(engine.sync_engine, "connect")
|
||||
def _set_pragmas(dbapi_connection: Any, _connection_record: Any) -> None:
|
||||
cursor = dbapi_connection.cursor()
|
||||
cursor.execute("PRAGMA journal_mode=WAL")
|
||||
cursor.execute("PRAGMA foreign_keys=ON")
|
||||
cursor.execute(f"PRAGMA busy_timeout={busy_timeout_ms}")
|
||||
cursor.close()
|
||||
|
||||
|
||||
__all__ = ["JobRow", "configure_sqlite_engine"]
|
||||
14
dify-agent/src/shellctl/server/errors.py
Normal file
14
dify-agent/src/shellctl/server/errors.py
Normal file
@ -0,0 +1,14 @@
|
||||
"""Shared exception types for shellctl server-side modules."""
|
||||
|
||||
|
||||
class ShellctlServerError(RuntimeError):
|
||||
"""Structured server-side error that maps directly to API responses."""
|
||||
|
||||
def __init__(self, status_code: int, code: str, message: str) -> None:
|
||||
super().__init__(message)
|
||||
self.status_code = status_code
|
||||
self.code = code
|
||||
self.message = message
|
||||
|
||||
|
||||
__all__ = ["ShellctlServerError"]
|
||||
103
dify-agent/src/shellctl/server/serve.py
Normal file
103
dify-agent/src/shellctl/server/serve.py
Normal file
@ -0,0 +1,103 @@
|
||||
"""Server-side `shellctl serve` implementation.
|
||||
|
||||
This module stays separate from the top-level CLI so ordinary client commands
|
||||
do not import the FastAPI or uvicorn stack. `serve_command()` performs those
|
||||
imports lazily because only the long-running server path needs them.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import typer
|
||||
|
||||
from shellctl.shared.constants import (
|
||||
DEFAULT_AUTH_TOKEN_ENV,
|
||||
DEFAULT_GC_FINISHED_JOB_RETENTION_SECONDS,
|
||||
DEFAULT_GC_INTERVAL_SECONDS,
|
||||
)
|
||||
from shellctl.shared.runtime import (
|
||||
default_state_dir,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from shellctl.server.config import ShellctlConfig
|
||||
|
||||
|
||||
def serve_command(
|
||||
listen: str = "127.0.0.1:8765",
|
||||
auth_token: str | None = typer.Option(
|
||||
None,
|
||||
"--auth-token",
|
||||
envvar=DEFAULT_AUTH_TOKEN_ENV,
|
||||
help=(
|
||||
"Bearer token value. You can also set SHELLCTL_AUTH_TOKEN. "
|
||||
"Leave it unset or empty to disable HTTP bearer auth."
|
||||
),
|
||||
),
|
||||
state_dir: Path | None = None,
|
||||
runtime_dir: Path | None = None,
|
||||
gc_interval_seconds: float = DEFAULT_GC_INTERVAL_SECONDS,
|
||||
gc_finished_job_retention_seconds: float = DEFAULT_GC_FINISHED_JOB_RETENTION_SECONDS,
|
||||
) -> None:
|
||||
"""Build `ShellctlConfig` from CLI inputs and run the local HTTP server.
|
||||
|
||||
Args:
|
||||
listen: Host/port pair for the uvicorn listener in `host:port` form.
|
||||
auth_token: Optional bearer token value. An explicit empty string
|
||||
disables HTTP auth, an explicit non-empty token enables it, and an
|
||||
omitted/`None` value may still resolve from `SHELLCTL_AUTH_TOKEN`
|
||||
through the Typer env var binding or `ShellctlConfig` fallback.
|
||||
state_dir: Persistent shellctl state directory; defaults to the shared
|
||||
XDG-style state path when omitted.
|
||||
runtime_dir: Optional runtime directory override for tmux/runtime
|
||||
artifacts.
|
||||
gc_interval_seconds: Background GC wake-up cadence for finished jobs.
|
||||
gc_finished_job_retention_seconds: Retention window before finished jobs
|
||||
are eligible for GC.
|
||||
|
||||
This entrypoint is the only CLI path that should pull in the FastAPI and
|
||||
uvicorn stack. It parses the listener, constructs `ShellctlConfig`, and
|
||||
then hands the configured app to uvicorn.
|
||||
"""
|
||||
|
||||
from shellctl.server.config import ShellctlConfig
|
||||
|
||||
host, port = _parse_listen(listen)
|
||||
config = ShellctlConfig(
|
||||
listen=listen,
|
||||
auth_token=auth_token,
|
||||
state_dir=state_dir or default_state_dir(),
|
||||
runtime_dir=runtime_dir,
|
||||
gc_interval_seconds=gc_interval_seconds,
|
||||
gc_finished_job_retention_seconds=gc_finished_job_retention_seconds,
|
||||
)
|
||||
_uvicorn_run(_create_app(config), host=host, port=port, log_level="info")
|
||||
|
||||
|
||||
def _parse_listen(value: str) -> tuple[str, int]:
|
||||
if ":" not in value:
|
||||
raise typer.BadParameter("listen must use host:port format")
|
||||
host, raw_port = value.rsplit(":", 1)
|
||||
host = host.strip("[]")
|
||||
try:
|
||||
port = int(raw_port)
|
||||
except ValueError as exc:
|
||||
raise typer.BadParameter(f"invalid port: {raw_port}") from exc
|
||||
return host, port
|
||||
|
||||
|
||||
def _create_app(config: ShellctlConfig) -> Any:
|
||||
from shellctl.server.api import create_app
|
||||
|
||||
return create_app(config)
|
||||
|
||||
|
||||
def _uvicorn_run(app: Any, *, host: str, port: int, log_level: str) -> None:
|
||||
import uvicorn
|
||||
|
||||
uvicorn.run(app, host=host, port=port, log_level=log_level)
|
||||
|
||||
|
||||
__all__ = ["serve_command"]
|
||||
1184
dify-agent/src/shellctl/server/service.py
Normal file
1184
dify-agent/src/shellctl/server/service.py
Normal file
File diff suppressed because it is too large
Load Diff
343
dify-agent/src/shellctl/server/tmux.py
Normal file
343
dify-agent/src/shellctl/server/tmux.py
Normal file
@ -0,0 +1,343 @@
|
||||
"""tmux control layer for shellctl jobs.
|
||||
|
||||
The rest of shellctl treats this module as the only place that knows tmux CLI
|
||||
command shapes. Tests can replace `TmuxControllerProtocol` with a fake to
|
||||
exercise SQLite/output semantics without depending on a local tmux daemon.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import shlex
|
||||
import subprocess
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import Protocol, cast
|
||||
|
||||
import anyio
|
||||
|
||||
from shellctl.server.artifacts import (
|
||||
pipe_drained_path,
|
||||
pipe_error_log_path,
|
||||
pipe_failed_path,
|
||||
runner_ended_at_path,
|
||||
runner_exit_code_path,
|
||||
)
|
||||
from shellctl.server.config import ShellctlConfig
|
||||
from shellctl.server.errors import ShellctlServerError
|
||||
from shellctl.shared.runtime import (
|
||||
job_pane_target,
|
||||
job_session_name,
|
||||
)
|
||||
from shellctl.shared.schemas import TerminalSize
|
||||
|
||||
|
||||
class TmuxControllerProtocol(Protocol):
|
||||
"""Protocol used by `ShellctlService` for tmux interactions."""
|
||||
|
||||
async def start_server(self) -> None: ...
|
||||
|
||||
async def list_sessions(self) -> set[str]: ...
|
||||
|
||||
async def session_exists(self, session_name: str) -> bool: ...
|
||||
|
||||
async def is_output_pipe_active(self, *, job_id: str) -> bool | None: ...
|
||||
|
||||
async def create_job_session(
|
||||
self,
|
||||
*,
|
||||
job_id: str,
|
||||
job_dir: Path,
|
||||
cwd: Path,
|
||||
terminal: TerminalSize,
|
||||
) -> None: ...
|
||||
|
||||
async def enable_output_pipe(self, *, job_id: str, job_dir: Path, ready_file: Path) -> None: ...
|
||||
|
||||
async def send_input(self, *, job_id: str, text: str) -> None: ...
|
||||
|
||||
async def send_interrupt(self, *, job_id: str) -> None: ...
|
||||
|
||||
async def cleanup_session(self, *, job_id: str) -> None: ...
|
||||
|
||||
|
||||
class TmuxController:
|
||||
"""Best-effort wrapper around a dedicated tmux socket.
|
||||
|
||||
The controller always clears `TMUX` from the child environment and always
|
||||
passes `-S <socket>` so shellctl sessions stay isolated from the user's
|
||||
default tmux server.
|
||||
"""
|
||||
|
||||
def __init__(self, config: ShellctlConfig) -> None:
|
||||
self._config = config
|
||||
|
||||
async def start_server(self) -> None:
|
||||
await self._run_tmux("start-server")
|
||||
|
||||
async def list_sessions(self) -> set[str]:
|
||||
result = await self._run_tmux("list-sessions", "-F", "#{session_name}", check=False)
|
||||
if result.returncode != 0:
|
||||
stderr = result.stderr.decode("utf-8", errors="replace")
|
||||
if _tmux_target_missing(stderr):
|
||||
return set()
|
||||
raise ShellctlServerError(500, "tmux_error", stderr.strip() or "tmux list-sessions failed")
|
||||
output = result.stdout.decode("utf-8", errors="replace")
|
||||
return {line.strip() for line in output.splitlines() if line.strip()}
|
||||
|
||||
async def session_exists(self, session_name: str) -> bool:
|
||||
return session_name in await self.list_sessions()
|
||||
|
||||
async def is_output_pipe_active(self, *, job_id: str) -> bool | None:
|
||||
result = await self._run_tmux(
|
||||
"display-message",
|
||||
"-p",
|
||||
"-t",
|
||||
job_pane_target(job_id),
|
||||
"#{pane_pipe}",
|
||||
check=False,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
stderr = result.stderr.decode("utf-8", errors="replace")
|
||||
if _tmux_target_missing(stderr):
|
||||
return None
|
||||
raise ShellctlServerError(
|
||||
500,
|
||||
"tmux_error",
|
||||
stderr.strip() or f"Failed to inspect output pipe for {job_id}",
|
||||
)
|
||||
return result.stdout.decode("utf-8", errors="replace").strip() == "1"
|
||||
|
||||
async def create_job_session(
|
||||
self,
|
||||
*,
|
||||
job_id: str,
|
||||
job_dir: Path,
|
||||
cwd: Path,
|
||||
terminal: TerminalSize,
|
||||
) -> None:
|
||||
runner_command = self._shell_join(
|
||||
[
|
||||
str(self._config.runner_path),
|
||||
str(job_dir),
|
||||
job_id,
|
||||
str(cwd),
|
||||
]
|
||||
)
|
||||
result = await self._run_tmux(
|
||||
"-f",
|
||||
"/dev/null",
|
||||
"new-session",
|
||||
"-d",
|
||||
"-s",
|
||||
job_session_name(job_id),
|
||||
"-x",
|
||||
str(terminal.cols),
|
||||
"-y",
|
||||
str(terminal.rows),
|
||||
runner_command,
|
||||
check=False,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
raise ShellctlServerError(
|
||||
500,
|
||||
"tmux_new_session_failed",
|
||||
result.stderr.decode("utf-8", errors="replace").strip()
|
||||
or f"Failed to create tmux session for {job_id}",
|
||||
)
|
||||
|
||||
async def enable_output_pipe(self, *, job_id: str, job_dir: Path, ready_file: Path) -> None:
|
||||
output_command = self._pipe_command_source(
|
||||
job_id=job_id,
|
||||
job_dir=job_dir,
|
||||
ready_file=ready_file,
|
||||
)
|
||||
result = await self._run_tmux(
|
||||
"pipe-pane",
|
||||
"-o",
|
||||
"-t",
|
||||
job_pane_target(job_id),
|
||||
output_command,
|
||||
check=False,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
raise ShellctlServerError(
|
||||
500,
|
||||
"pipe_failed",
|
||||
result.stderr.decode("utf-8", errors="replace").strip() or f"Failed to attach output pipe for {job_id}",
|
||||
)
|
||||
|
||||
def _pipe_command_source(self, *, job_id: str, job_dir: Path, ready_file: Path) -> str:
|
||||
"""Build the tmux `pipe-pane` command that drains and finalizes output.
|
||||
|
||||
For normal exits, the runner now records completion metadata into job
|
||||
artifacts and the pipe finalizer commits `runner-exit` only after
|
||||
the lightweight sanitizer reaches EOF and flushes `output.log`
|
||||
successfully. Sanitizer stderr is captured into `pipe-error.log` so
|
||||
startup timeouts can distinguish slow imports from subprocess crashes.
|
||||
If the follow-up `runner-exit` write fails, the drain marker remains in
|
||||
place and stderr is appended to the same log with an explicit status
|
||||
line. The pipe still exits with the sanitizer status so a drained job is
|
||||
not misclassified as `pipe_failed` before reconciliation can recover the
|
||||
SQLite write from the drained artifacts.
|
||||
"""
|
||||
|
||||
sanitize_command = self._shell_join(
|
||||
(
|
||||
*self._config.sanitize_pty_command,
|
||||
"--ready-file",
|
||||
str(ready_file),
|
||||
)
|
||||
)
|
||||
runner_exit_command = self._shell_join(
|
||||
(
|
||||
*self._config.runner_exit_command,
|
||||
"--state-dir",
|
||||
str(self._config.state_dir),
|
||||
"--job-id",
|
||||
job_id,
|
||||
"--sqlite-busy-timeout-ms",
|
||||
str(self._config.sqlite_busy_timeout_ms),
|
||||
)
|
||||
)
|
||||
output_path = shlex.quote(str(job_dir / "output.log"))
|
||||
drained_path = shlex.quote(str(pipe_drained_path(job_dir)))
|
||||
error_log_path = shlex.quote(str(pipe_error_log_path(job_dir)))
|
||||
failed_path = shlex.quote(str(pipe_failed_path(job_dir)))
|
||||
exit_code_path = shlex.quote(str(runner_exit_code_path(job_dir)))
|
||||
ended_at_path = shlex.quote(str(runner_ended_at_path(job_dir)))
|
||||
return " ; ".join(
|
||||
[
|
||||
f"{sanitize_command} >> {output_path} 2> {error_log_path}",
|
||||
"sanitize_status=$?",
|
||||
"runner_exit_status=0",
|
||||
(
|
||||
'if [ "$sanitize_status" -eq 0 ]; then '
|
||||
f": > {drained_path}; "
|
||||
f"if [ -s {exit_code_path} ] && [ -s {ended_at_path} ]; then "
|
||||
f'{runner_exit_command} --exit-code "$(cat {exit_code_path})" '
|
||||
f'--ended-at "$(cat {ended_at_path})" 2>> {error_log_path}; '
|
||||
"runner_exit_status=$?; "
|
||||
'if [ "$runner_exit_status" -ne 0 ]; then '
|
||||
f"printf 'runner-exit failed with status %s\\n' \"$runner_exit_status\" >> {error_log_path}; "
|
||||
"fi; fi; "
|
||||
f"else : > {failed_path}; fi"
|
||||
),
|
||||
'if [ "$sanitize_status" -ne 0 ]; then exit "$sanitize_status"; fi',
|
||||
'exit "$sanitize_status"',
|
||||
]
|
||||
)
|
||||
|
||||
async def send_input(self, *, job_id: str, text: str) -> None:
|
||||
buffer_name = f"shellctl-in-{job_id}"
|
||||
runtime_dir = cast(Path, self._config.runtime_dir)
|
||||
runtime_dir.mkdir(parents=True, exist_ok=True)
|
||||
fd, tmp_name = tempfile.mkstemp(prefix=f"shellctl-input-{job_id}-", dir=runtime_dir)
|
||||
tmp_path = Path(tmp_name)
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as handle:
|
||||
handle.write(text)
|
||||
load_result = await self._run_tmux(
|
||||
"load-buffer",
|
||||
"-b",
|
||||
buffer_name,
|
||||
str(tmp_path),
|
||||
check=False,
|
||||
)
|
||||
if load_result.returncode != 0:
|
||||
stderr = load_result.stderr.decode("utf-8", errors="replace").strip()
|
||||
if _tmux_target_missing(stderr):
|
||||
raise ShellctlServerError(
|
||||
409,
|
||||
"tmux_target_missing",
|
||||
stderr or f"The tmux pane for {job_id} is no longer available",
|
||||
)
|
||||
raise ShellctlServerError(
|
||||
500,
|
||||
"tmux_input_failed",
|
||||
stderr or f"Failed to load input buffer for {job_id}",
|
||||
)
|
||||
paste_result = await self._run_tmux(
|
||||
"paste-buffer",
|
||||
"-t",
|
||||
job_pane_target(job_id),
|
||||
"-b",
|
||||
buffer_name,
|
||||
check=False,
|
||||
)
|
||||
if paste_result.returncode != 0:
|
||||
stderr = paste_result.stderr.decode("utf-8", errors="replace").strip()
|
||||
if _tmux_target_missing(stderr):
|
||||
raise ShellctlServerError(
|
||||
409,
|
||||
"tmux_target_missing",
|
||||
stderr or f"The tmux pane for {job_id} is no longer available",
|
||||
)
|
||||
raise ShellctlServerError(
|
||||
500,
|
||||
"tmux_input_failed",
|
||||
stderr or f"Failed to paste input buffer for {job_id}",
|
||||
)
|
||||
finally:
|
||||
await self._run_tmux("delete-buffer", "-b", buffer_name, check=False)
|
||||
tmp_path.unlink(missing_ok=True)
|
||||
|
||||
async def send_interrupt(self, *, job_id: str) -> None:
|
||||
await self._run_tmux(
|
||||
"send-keys",
|
||||
"-t",
|
||||
job_pane_target(job_id),
|
||||
"C-c",
|
||||
check=False,
|
||||
)
|
||||
|
||||
async def cleanup_session(self, *, job_id: str) -> None:
|
||||
await self._run_tmux(
|
||||
"kill-session",
|
||||
"-t",
|
||||
job_session_name(job_id),
|
||||
check=False,
|
||||
)
|
||||
|
||||
async def _run_tmux(self, *args: str, check: bool = True) -> subprocess.CompletedProcess[bytes]:
|
||||
env = dict(os.environ)
|
||||
env.pop("TMUX", None)
|
||||
try:
|
||||
result = await anyio.run_process(
|
||||
["tmux", "-S", str(self._config.tmux_socket), *args],
|
||||
env=env,
|
||||
check=False,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
)
|
||||
except FileNotFoundError as exc:
|
||||
raise ShellctlServerError(500, "tmux_not_installed", "tmux executable was not found") from exc
|
||||
if check and result.returncode != 0:
|
||||
raise ShellctlServerError(
|
||||
500,
|
||||
"tmux_error",
|
||||
result.stderr.decode("utf-8", errors="replace").strip() or "tmux command failed",
|
||||
)
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _shell_join(parts: tuple[str, ...] | list[str]) -> str:
|
||||
return " ".join(shlex.quote(part) for part in parts)
|
||||
|
||||
|
||||
def _tmux_target_missing(stderr: str) -> bool:
|
||||
normalized = stderr.lower()
|
||||
return (
|
||||
"can't find pane" in normalized
|
||||
or "can't find session" in normalized
|
||||
or "no server running" in normalized
|
||||
or "failed to connect" in normalized
|
||||
or "server exited unexpectedly" in normalized
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"TmuxController",
|
||||
"TmuxControllerProtocol",
|
||||
"_tmux_target_missing",
|
||||
]
|
||||
182
dify-agent/src/shellctl/shared/__init__.py
Normal file
182
dify-agent/src/shellctl/shared/__init__.py
Normal file
@ -0,0 +1,182 @@
|
||||
"""Shared shellctl transport/runtime helpers.
|
||||
|
||||
This package preserves the historical import surface while keeping the package
|
||||
root lazy. Lightweight callers can import concrete submodules such as
|
||||
`shared.runtime` without eagerly importing the pydantic schema layer, output
|
||||
or helpers outside the shared compatibility surface.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from importlib import import_module
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from shellctl.shared.constants import (
|
||||
DEFAULT_AUTH_TOKEN_ENV,
|
||||
DEFAULT_BASE_URL,
|
||||
DEFAULT_BASE_URL_ENV,
|
||||
DEFAULT_GC_FINISHED_JOB_RETENTION_SECONDS,
|
||||
DEFAULT_GC_INTERVAL_SECONDS,
|
||||
DEFAULT_HEALTH_STATUS,
|
||||
DEFAULT_IDLE_FLUSH_SECONDS,
|
||||
DEFAULT_LIST_LIMIT,
|
||||
DEFAULT_OUTPUT_LIMIT_BYTES,
|
||||
DEFAULT_TERMINAL_COLS,
|
||||
DEFAULT_TERMINAL_ROWS,
|
||||
DEFAULT_TERMINATE_GRACE_SECONDS,
|
||||
DEFAULT_TIMEOUT_SECONDS,
|
||||
JOB_ID_ALPHABET,
|
||||
JOB_ID_RANDOM_SUFFIX_LENGTH,
|
||||
MAX_LIST_LIMIT,
|
||||
MAX_OUTPUT_LIMIT_BYTES,
|
||||
MAX_WAIT_TIMEOUT_SECONDS,
|
||||
SESSION_NAME_PREFIX,
|
||||
)
|
||||
from shellctl.shared.output import (
|
||||
OutputWindow,
|
||||
read_output_window,
|
||||
tail_output_window,
|
||||
)
|
||||
from shellctl.shared.runtime import (
|
||||
default_runtime_dir,
|
||||
default_state_dir,
|
||||
format_timestamp,
|
||||
generate_job_id,
|
||||
is_terminal_status,
|
||||
job_pane_target,
|
||||
job_session_name,
|
||||
parse_timestamp,
|
||||
utc_now,
|
||||
)
|
||||
from shellctl.shared.schemas import (
|
||||
TERMINAL_JOB_STATUSES,
|
||||
DeleteJobResponse,
|
||||
ErrorDetail,
|
||||
ErrorResponse,
|
||||
HealthResponse,
|
||||
InputJobRequest,
|
||||
JobInfo,
|
||||
JobResult,
|
||||
JobStatusName,
|
||||
JobStatusView,
|
||||
ListJobsResponse,
|
||||
RunJobRequest,
|
||||
ShellctlModel,
|
||||
TerminalSize,
|
||||
TerminateJobRequest,
|
||||
WaitJobRequest,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"DEFAULT_AUTH_TOKEN_ENV",
|
||||
"DEFAULT_BASE_URL",
|
||||
"DEFAULT_BASE_URL_ENV",
|
||||
"DEFAULT_GC_FINISHED_JOB_RETENTION_SECONDS",
|
||||
"DEFAULT_GC_INTERVAL_SECONDS",
|
||||
"DEFAULT_HEALTH_STATUS",
|
||||
"DEFAULT_IDLE_FLUSH_SECONDS",
|
||||
"DEFAULT_LIST_LIMIT",
|
||||
"DEFAULT_OUTPUT_LIMIT_BYTES",
|
||||
"DEFAULT_TERMINAL_COLS",
|
||||
"DEFAULT_TERMINAL_ROWS",
|
||||
"DEFAULT_TERMINATE_GRACE_SECONDS",
|
||||
"DEFAULT_TIMEOUT_SECONDS",
|
||||
"JOB_ID_ALPHABET",
|
||||
"JOB_ID_RANDOM_SUFFIX_LENGTH",
|
||||
"MAX_LIST_LIMIT",
|
||||
"MAX_OUTPUT_LIMIT_BYTES",
|
||||
"MAX_WAIT_TIMEOUT_SECONDS",
|
||||
"SESSION_NAME_PREFIX",
|
||||
"TERMINAL_JOB_STATUSES",
|
||||
"DeleteJobResponse",
|
||||
"ErrorDetail",
|
||||
"ErrorResponse",
|
||||
"HealthResponse",
|
||||
"InputJobRequest",
|
||||
"JobInfo",
|
||||
"JobResult",
|
||||
"JobStatusName",
|
||||
"JobStatusView",
|
||||
"ListJobsResponse",
|
||||
"OutputWindow",
|
||||
"RunJobRequest",
|
||||
"ShellctlModel",
|
||||
"TerminalSize",
|
||||
"TerminateJobRequest",
|
||||
"WaitJobRequest",
|
||||
"default_runtime_dir",
|
||||
"default_state_dir",
|
||||
"format_timestamp",
|
||||
"generate_job_id",
|
||||
"is_terminal_status",
|
||||
"job_pane_target",
|
||||
"job_session_name",
|
||||
"parse_timestamp",
|
||||
"read_output_window",
|
||||
"tail_output_window",
|
||||
"utc_now",
|
||||
]
|
||||
|
||||
_EXPORTS = {
|
||||
"DEFAULT_AUTH_TOKEN_ENV": "shellctl.shared.constants",
|
||||
"DEFAULT_BASE_URL": "shellctl.shared.constants",
|
||||
"DEFAULT_BASE_URL_ENV": "shellctl.shared.constants",
|
||||
"DEFAULT_GC_FINISHED_JOB_RETENTION_SECONDS": "shellctl.shared.constants",
|
||||
"DEFAULT_GC_INTERVAL_SECONDS": "shellctl.shared.constants",
|
||||
"DEFAULT_HEALTH_STATUS": "shellctl.shared.constants",
|
||||
"DEFAULT_IDLE_FLUSH_SECONDS": "shellctl.shared.constants",
|
||||
"DEFAULT_LIST_LIMIT": "shellctl.shared.constants",
|
||||
"DEFAULT_OUTPUT_LIMIT_BYTES": "shellctl.shared.constants",
|
||||
"DEFAULT_TERMINAL_COLS": "shellctl.shared.constants",
|
||||
"DEFAULT_TERMINAL_ROWS": "shellctl.shared.constants",
|
||||
"DEFAULT_TERMINATE_GRACE_SECONDS": "shellctl.shared.constants",
|
||||
"DEFAULT_TIMEOUT_SECONDS": "shellctl.shared.constants",
|
||||
"JOB_ID_ALPHABET": "shellctl.shared.constants",
|
||||
"JOB_ID_RANDOM_SUFFIX_LENGTH": "shellctl.shared.constants",
|
||||
"MAX_LIST_LIMIT": "shellctl.shared.constants",
|
||||
"MAX_OUTPUT_LIMIT_BYTES": "shellctl.shared.constants",
|
||||
"MAX_WAIT_TIMEOUT_SECONDS": "shellctl.shared.constants",
|
||||
"SESSION_NAME_PREFIX": "shellctl.shared.constants",
|
||||
"OutputWindow": "shellctl.shared.output",
|
||||
"read_output_window": "shellctl.shared.output",
|
||||
"tail_output_window": "shellctl.shared.output",
|
||||
"default_runtime_dir": "shellctl.shared.runtime",
|
||||
"default_state_dir": "shellctl.shared.runtime",
|
||||
"format_timestamp": "shellctl.shared.runtime",
|
||||
"generate_job_id": "shellctl.shared.runtime",
|
||||
"is_terminal_status": "shellctl.shared.runtime",
|
||||
"job_pane_target": "shellctl.shared.runtime",
|
||||
"job_session_name": "shellctl.shared.runtime",
|
||||
"parse_timestamp": "shellctl.shared.runtime",
|
||||
"utc_now": "shellctl.shared.runtime",
|
||||
"TERMINAL_JOB_STATUSES": "shellctl.shared.schemas",
|
||||
"DeleteJobResponse": "shellctl.shared.schemas",
|
||||
"ErrorDetail": "shellctl.shared.schemas",
|
||||
"ErrorResponse": "shellctl.shared.schemas",
|
||||
"HealthResponse": "shellctl.shared.schemas",
|
||||
"InputJobRequest": "shellctl.shared.schemas",
|
||||
"JobInfo": "shellctl.shared.schemas",
|
||||
"JobResult": "shellctl.shared.schemas",
|
||||
"JobStatusName": "shellctl.shared.schemas",
|
||||
"JobStatusView": "shellctl.shared.schemas",
|
||||
"ListJobsResponse": "shellctl.shared.schemas",
|
||||
"RunJobRequest": "shellctl.shared.schemas",
|
||||
"ShellctlModel": "shellctl.shared.schemas",
|
||||
"TerminalSize": "shellctl.shared.schemas",
|
||||
"TerminateJobRequest": "shellctl.shared.schemas",
|
||||
"WaitJobRequest": "shellctl.shared.schemas",
|
||||
}
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
if name not in _EXPORTS:
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
module = import_module(_EXPORTS[name])
|
||||
value = getattr(module, name) # noqa: no-new-getattr lazy export proxy
|
||||
globals()[name] = value
|
||||
return value
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return sorted(set(globals()) | set(__all__))
|
||||
49
dify-agent/src/shellctl/shared/constants.py
Normal file
49
dify-agent/src/shellctl/shared/constants.py
Normal file
@ -0,0 +1,49 @@
|
||||
"""Shared shellctl constants.
|
||||
|
||||
This module intentionally contains only literal defaults and bounds so callers
|
||||
can import stable configuration values without pulling in the heavier DTO,
|
||||
sanitize, or server packages. The network CLI depends on this file for its
|
||||
base URL and auth-token environment contract, so keep it import-light.
|
||||
"""
|
||||
|
||||
DEFAULT_AUTH_TOKEN_ENV = "SHELLCTL_AUTH_TOKEN"
|
||||
DEFAULT_BASE_URL_ENV = "SHELLCTL_BASE_URL"
|
||||
DEFAULT_BASE_URL = "http://127.0.0.1:8765"
|
||||
DEFAULT_OUTPUT_LIMIT_BYTES = 1024 * 8
|
||||
MAX_OUTPUT_LIMIT_BYTES = 1024 * 1024
|
||||
DEFAULT_TIMEOUT_SECONDS = 30.0
|
||||
MAX_WAIT_TIMEOUT_SECONDS = 5.0 * 60.0
|
||||
DEFAULT_IDLE_FLUSH_SECONDS = 0.5
|
||||
DEFAULT_TERMINATE_GRACE_SECONDS = 5.0
|
||||
DEFAULT_TERMINAL_COLS = 120
|
||||
DEFAULT_TERMINAL_ROWS = 80
|
||||
DEFAULT_LIST_LIMIT = 100
|
||||
MAX_LIST_LIMIT = 1000
|
||||
DEFAULT_GC_INTERVAL_SECONDS = 10.0 * 60.0
|
||||
DEFAULT_GC_FINISHED_JOB_RETENTION_SECONDS = 24.0 * 60.0 * 60.0
|
||||
JOB_ID_ALPHABET = "0123456789abcdefghjkmnpqrstvwxyz"
|
||||
JOB_ID_RANDOM_SUFFIX_LENGTH = 3
|
||||
SESSION_NAME_PREFIX = "shellctl-job-"
|
||||
DEFAULT_HEALTH_STATUS = "ok"
|
||||
|
||||
__all__ = [
|
||||
"DEFAULT_AUTH_TOKEN_ENV",
|
||||
"DEFAULT_BASE_URL",
|
||||
"DEFAULT_BASE_URL_ENV",
|
||||
"DEFAULT_GC_FINISHED_JOB_RETENTION_SECONDS",
|
||||
"DEFAULT_GC_INTERVAL_SECONDS",
|
||||
"DEFAULT_HEALTH_STATUS",
|
||||
"DEFAULT_IDLE_FLUSH_SECONDS",
|
||||
"DEFAULT_LIST_LIMIT",
|
||||
"DEFAULT_OUTPUT_LIMIT_BYTES",
|
||||
"DEFAULT_TERMINAL_COLS",
|
||||
"DEFAULT_TERMINAL_ROWS",
|
||||
"DEFAULT_TERMINATE_GRACE_SECONDS",
|
||||
"DEFAULT_TIMEOUT_SECONDS",
|
||||
"JOB_ID_ALPHABET",
|
||||
"JOB_ID_RANDOM_SUFFIX_LENGTH",
|
||||
"MAX_LIST_LIMIT",
|
||||
"MAX_OUTPUT_LIMIT_BYTES",
|
||||
"MAX_WAIT_TIMEOUT_SECONDS",
|
||||
"SESSION_NAME_PREFIX",
|
||||
]
|
||||
120
dify-agent/src/shellctl/shared/output.py
Normal file
120
dify-agent/src/shellctl/shared/output.py
Normal file
@ -0,0 +1,120 @@
|
||||
"""UTF-8-safe output slicing helpers for shellctl."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
@dataclass(slots=True, frozen=True)
|
||||
class OutputWindow:
|
||||
"""UTF-8-safe slice of `output.log`."""
|
||||
|
||||
output: str
|
||||
offset: int
|
||||
truncated: bool
|
||||
|
||||
|
||||
def read_output_window(path: Path, *, offset: int, limit: int) -> OutputWindow:
|
||||
"""Read a forward UTF-8-safe slice from `output.log`."""
|
||||
|
||||
if offset < 0:
|
||||
raise ValueError(f"offset must be >= 0; got {offset}")
|
||||
if limit <= 0:
|
||||
raise ValueError(f"limit must be > 0; got {limit}")
|
||||
|
||||
if not path.exists():
|
||||
if offset == 0:
|
||||
return OutputWindow(output="", offset=0, truncated=False)
|
||||
raise ValueError(f"offset {offset} exceeds current file size 0")
|
||||
|
||||
size = path.stat().st_size
|
||||
if offset > size:
|
||||
raise ValueError(f"offset {offset} exceeds current file size {size}")
|
||||
if offset == size:
|
||||
return OutputWindow(output="", offset=offset, truncated=False)
|
||||
|
||||
with path.open("rb") as handle:
|
||||
handle.seek(offset)
|
||||
raw = handle.read(limit + 4)
|
||||
|
||||
start_shift = _advance_to_utf8_boundary(raw, 0)
|
||||
data = raw[start_shift:]
|
||||
budget = max(0, limit - start_shift)
|
||||
consumed = _valid_utf8_prefix_length(data[:budget])
|
||||
if consumed == 0 and data:
|
||||
consumed = _first_complete_utf8_char_length(data)
|
||||
output_bytes = data[:consumed]
|
||||
new_offset = offset + start_shift + consumed
|
||||
truncated = new_offset < size
|
||||
return OutputWindow(
|
||||
output=output_bytes.decode("utf-8", errors="strict"),
|
||||
offset=new_offset,
|
||||
truncated=truncated,
|
||||
)
|
||||
|
||||
|
||||
def tail_output_window(path: Path, *, limit: int) -> OutputWindow:
|
||||
"""Read a UTF-8-safe tail snapshot from `output.log`."""
|
||||
|
||||
if limit <= 0:
|
||||
raise ValueError(f"limit must be > 0; got {limit}")
|
||||
if not path.exists():
|
||||
return OutputWindow(output="", offset=0, truncated=False)
|
||||
size = path.stat().st_size
|
||||
if size == 0:
|
||||
return OutputWindow(output="", offset=0, truncated=False)
|
||||
start = max(0, size - limit)
|
||||
padded_start = max(0, start - 4)
|
||||
with path.open("rb") as handle:
|
||||
handle.seek(padded_start)
|
||||
raw = handle.read(size - padded_start)
|
||||
relative_start = _advance_to_utf8_boundary(raw, start - padded_start)
|
||||
payload = raw[relative_start:]
|
||||
consumed = _valid_utf8_prefix_length(payload)
|
||||
output_bytes = payload[:consumed]
|
||||
return OutputWindow(
|
||||
output=output_bytes.decode("utf-8", errors="strict"),
|
||||
offset=padded_start + relative_start + consumed,
|
||||
truncated=False,
|
||||
)
|
||||
|
||||
|
||||
def _advance_to_utf8_boundary(data: bytes, start: int) -> int:
|
||||
while start < len(data) and _is_utf8_continuation_byte(data[start]):
|
||||
start += 1
|
||||
return start
|
||||
|
||||
|
||||
def _valid_utf8_prefix_length(data: bytes) -> int:
|
||||
end = len(data)
|
||||
while end >= 0:
|
||||
try:
|
||||
data[:end].decode("utf-8", errors="strict")
|
||||
return end
|
||||
except UnicodeDecodeError as exc:
|
||||
if exc.start < end - 4:
|
||||
end = exc.start
|
||||
else:
|
||||
end -= 1
|
||||
return 0
|
||||
|
||||
|
||||
def _is_utf8_continuation_byte(value: int) -> bool:
|
||||
return (value & 0b1100_0000) == 0b1000_0000
|
||||
|
||||
|
||||
def _first_complete_utf8_char_length(data: bytes) -> int:
|
||||
for end in range(1, len(data) + 1):
|
||||
try:
|
||||
decoded = data[:end].decode("utf-8", errors="strict")
|
||||
except UnicodeDecodeError:
|
||||
continue
|
||||
if len(decoded) == 1:
|
||||
return end
|
||||
if decoded:
|
||||
return len(decoded[0].encode("utf-8"))
|
||||
return 0
|
||||
|
||||
|
||||
__all__ = ["OutputWindow", "read_output_window", "tail_output_window"]
|
||||
95
dify-agent/src/shellctl/shared/runtime.py
Normal file
95
dify-agent/src/shellctl/shared/runtime.py
Normal file
@ -0,0 +1,95 @@
|
||||
"""Stable runtime/naming helpers shared across shellctl modules."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import secrets
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
|
||||
from shellctl.shared.constants import (
|
||||
JOB_ID_ALPHABET,
|
||||
JOB_ID_RANDOM_SUFFIX_LENGTH,
|
||||
SESSION_NAME_PREFIX,
|
||||
)
|
||||
from shellctl.shared.schemas import (
|
||||
TERMINAL_JOB_STATUSES,
|
||||
JobStatusName,
|
||||
)
|
||||
|
||||
|
||||
def utc_now() -> datetime:
|
||||
"""Return the current UTC time with timezone information."""
|
||||
|
||||
return datetime.now(UTC)
|
||||
|
||||
|
||||
def format_timestamp(value: datetime | None = None) -> str:
|
||||
"""Format a UTC timestamp in the stable artifact/API representation."""
|
||||
|
||||
moment = (value or utc_now()).astimezone(UTC).replace(microsecond=0)
|
||||
return moment.isoformat().replace("+00:00", "Z")
|
||||
|
||||
|
||||
def parse_timestamp(value: str) -> datetime:
|
||||
"""Parse an artifact/API timestamp back into a timezone-aware datetime."""
|
||||
|
||||
return datetime.fromisoformat(value.replace("Z", "+00:00")).astimezone(UTC)
|
||||
|
||||
|
||||
def is_terminal_status(status: JobStatusName) -> bool:
|
||||
"""Check whether a lifecycle state is terminal."""
|
||||
|
||||
return status in TERMINAL_JOB_STATUSES
|
||||
|
||||
|
||||
def generate_job_id(*, now: datetime | None = None) -> str:
|
||||
"""Generate a short, human-readable job id."""
|
||||
|
||||
timestamp = (now or utc_now()).astimezone(UTC).strftime("%m%d%H%M")
|
||||
suffix = "".join(secrets.choice(JOB_ID_ALPHABET) for _ in range(JOB_ID_RANDOM_SUFFIX_LENGTH))
|
||||
return f"{timestamp}-{suffix}"
|
||||
|
||||
|
||||
def job_session_name(job_id: str) -> str:
|
||||
"""Return the dedicated tmux session name for a job."""
|
||||
|
||||
return f"{SESSION_NAME_PREFIX}{job_id}"
|
||||
|
||||
|
||||
def job_pane_target(job_id: str) -> str:
|
||||
"""Return the canonical tmux pane target for a single-pane job session."""
|
||||
|
||||
return f"{job_session_name(job_id)}:0.0"
|
||||
|
||||
|
||||
def default_state_dir() -> Path:
|
||||
"""Resolve the XDG-style default shellctl state directory."""
|
||||
|
||||
xdg_state_home = os.environ.get("XDG_STATE_HOME")
|
||||
if xdg_state_home:
|
||||
return Path(xdg_state_home) / "shellctl"
|
||||
return Path.home() / ".local" / "state" / "shellctl"
|
||||
|
||||
|
||||
def default_runtime_dir(state_dir: Path | None = None) -> Path:
|
||||
"""Resolve the XDG-style default shellctl runtime directory."""
|
||||
|
||||
xdg_runtime_dir = os.environ.get("XDG_RUNTIME_DIR")
|
||||
if xdg_runtime_dir:
|
||||
return Path(xdg_runtime_dir) / "shellctl"
|
||||
base_state_dir = state_dir or default_state_dir()
|
||||
return base_state_dir / "run" / "shellctl"
|
||||
|
||||
|
||||
__all__ = [
|
||||
"default_runtime_dir",
|
||||
"default_state_dir",
|
||||
"format_timestamp",
|
||||
"generate_job_id",
|
||||
"is_terminal_status",
|
||||
"job_pane_target",
|
||||
"job_session_name",
|
||||
"parse_timestamp",
|
||||
"utc_now",
|
||||
]
|
||||
213
dify-agent/src/shellctl/shared/schemas.py
Normal file
213
dify-agent/src/shellctl/shared/schemas.py
Normal file
@ -0,0 +1,213 @@
|
||||
"""Shared pydantic transport models for shellctl."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from enum import StrEnum
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator
|
||||
|
||||
from shellctl.shared.constants import (
|
||||
DEFAULT_HEALTH_STATUS,
|
||||
DEFAULT_IDLE_FLUSH_SECONDS,
|
||||
DEFAULT_OUTPUT_LIMIT_BYTES,
|
||||
DEFAULT_TERMINAL_COLS,
|
||||
DEFAULT_TERMINAL_ROWS,
|
||||
DEFAULT_TERMINATE_GRACE_SECONDS,
|
||||
DEFAULT_TIMEOUT_SECONDS,
|
||||
MAX_OUTPUT_LIMIT_BYTES,
|
||||
MAX_WAIT_TIMEOUT_SECONDS,
|
||||
)
|
||||
|
||||
|
||||
class ShellctlModel(BaseModel):
|
||||
"""Base pydantic model with strict extra-field handling.
|
||||
|
||||
shellctl uses these DTOs directly for HTTP request/response bodies, so
|
||||
silently accepting unknown fields would make it harder to detect schema
|
||||
drift.
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
|
||||
class JobStatusName(StrEnum):
|
||||
"""Lifecycle states materialized into SQLite rows and API responses."""
|
||||
|
||||
CREATED = "created"
|
||||
STARTING = "starting"
|
||||
RUNNING = "running"
|
||||
EXITED = "exited"
|
||||
TERMINATED = "terminated"
|
||||
FAILED = "failed"
|
||||
LOST = "lost"
|
||||
|
||||
|
||||
TERMINAL_JOB_STATUSES = frozenset(
|
||||
{
|
||||
JobStatusName.EXITED,
|
||||
JobStatusName.TERMINATED,
|
||||
JobStatusName.FAILED,
|
||||
JobStatusName.LOST,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class TerminalSize(ShellctlModel):
|
||||
"""Requested initial PTY geometry for a job."""
|
||||
|
||||
cols: int = Field(default=DEFAULT_TERMINAL_COLS, ge=1, le=4096)
|
||||
rows: int = Field(default=DEFAULT_TERMINAL_ROWS, ge=1, le=4096)
|
||||
|
||||
|
||||
class JobResult(ShellctlModel):
|
||||
"""Unified response shape for output-oriented job APIs."""
|
||||
|
||||
job_id: str
|
||||
done: bool
|
||||
status: JobStatusName
|
||||
exit_code: int | None = None
|
||||
output_path: str
|
||||
output: str
|
||||
offset: int = Field(ge=0)
|
||||
truncated: bool
|
||||
|
||||
|
||||
class JobStatusView(ShellctlModel):
|
||||
"""Materialized lifecycle view returned by status-like APIs."""
|
||||
|
||||
job_id: str
|
||||
status: JobStatusName
|
||||
done: bool
|
||||
exit_code: int | None = None
|
||||
created_at: str
|
||||
started_at: str | None = None
|
||||
ended_at: str | None = None
|
||||
offset: int = Field(ge=0)
|
||||
|
||||
|
||||
class JobInfo(ShellctlModel):
|
||||
"""Compact job listing record."""
|
||||
|
||||
job_id: str
|
||||
status: JobStatusName
|
||||
created_at: str
|
||||
started_at: str | None = None
|
||||
ended_at: str | None = None
|
||||
|
||||
|
||||
class ListJobsResponse(ShellctlModel):
|
||||
"""Response body for `GET /v1/jobs`."""
|
||||
|
||||
jobs: list[JobInfo]
|
||||
|
||||
|
||||
class DeleteJobResponse(ShellctlModel):
|
||||
"""Response body for successful delete operations."""
|
||||
|
||||
job_id: str
|
||||
deleted: bool = True
|
||||
|
||||
|
||||
class HealthResponse(ShellctlModel):
|
||||
"""Public health check response."""
|
||||
|
||||
status: str = DEFAULT_HEALTH_STATUS
|
||||
|
||||
|
||||
class ErrorDetail(ShellctlModel):
|
||||
"""Machine-readable API error payload."""
|
||||
|
||||
code: str
|
||||
message: str
|
||||
|
||||
|
||||
class ErrorResponse(ShellctlModel):
|
||||
"""Envelope used by server-side exception handlers."""
|
||||
|
||||
error: ErrorDetail
|
||||
|
||||
|
||||
class RunJobRequest(ShellctlModel):
|
||||
"""HTTP request body for `POST /v1/jobs/run`.
|
||||
|
||||
`env` augments the runner's inherited process environment instead of
|
||||
replacing it, so callers can preset script-local variables without losing
|
||||
ambient values such as `PATH`.
|
||||
"""
|
||||
|
||||
script: str
|
||||
cwd: str | None = None
|
||||
env: dict[str, str] | None = None
|
||||
terminal: TerminalSize | None = None
|
||||
timeout: float = Field(default=DEFAULT_TIMEOUT_SECONDS, gt=0, le=MAX_WAIT_TIMEOUT_SECONDS)
|
||||
output_limit: int = Field(default=DEFAULT_OUTPUT_LIMIT_BYTES, ge=1, le=MAX_OUTPUT_LIMIT_BYTES)
|
||||
idle_flush_seconds: float = Field(default=DEFAULT_IDLE_FLUSH_SECONDS, ge=0, le=30)
|
||||
|
||||
@field_validator("env")
|
||||
@classmethod
|
||||
def _validate_env(cls, env: dict[str, str] | None) -> dict[str, str] | None:
|
||||
"""Reject env entries that cannot be represented in `execve`.
|
||||
|
||||
shellctl applies `env` as a process environment overlay, so validation
|
||||
follows the low-level `NAME=value` constraints instead of shell variable
|
||||
naming rules: names must be non-empty and cannot contain `=` or NUL,
|
||||
while values cannot contain NUL.
|
||||
"""
|
||||
|
||||
if env is None:
|
||||
return None
|
||||
for name, value in env.items():
|
||||
if not name:
|
||||
raise ValueError("env names must be non-empty")
|
||||
if "=" in name:
|
||||
raise ValueError(f"env name must not contain '=': {name!r}")
|
||||
if "\x00" in name:
|
||||
raise ValueError(f"env name must not contain NUL: {name!r}")
|
||||
if "\x00" in value:
|
||||
raise ValueError(f"env value must not contain NUL: {name!r}")
|
||||
return env
|
||||
|
||||
|
||||
class WaitJobRequest(ShellctlModel):
|
||||
"""HTTP request body for `POST /v1/jobs/{job_id}/wait`."""
|
||||
|
||||
timeout: float = Field(default=DEFAULT_TIMEOUT_SECONDS, ge=0, le=MAX_WAIT_TIMEOUT_SECONDS)
|
||||
offset: int = Field(ge=0)
|
||||
output_limit: int = Field(default=DEFAULT_OUTPUT_LIMIT_BYTES, ge=1, le=MAX_OUTPUT_LIMIT_BYTES)
|
||||
idle_flush_seconds: float = Field(default=DEFAULT_IDLE_FLUSH_SECONDS, ge=0, le=30)
|
||||
|
||||
|
||||
class InputJobRequest(ShellctlModel):
|
||||
"""HTTP request body for `POST /v1/jobs/{job_id}/input`."""
|
||||
|
||||
text: str
|
||||
timeout: float = Field(default=DEFAULT_TIMEOUT_SECONDS, gt=0, le=MAX_WAIT_TIMEOUT_SECONDS)
|
||||
offset: int = Field(ge=0)
|
||||
output_limit: int = Field(default=DEFAULT_OUTPUT_LIMIT_BYTES, ge=1, le=MAX_OUTPUT_LIMIT_BYTES)
|
||||
idle_flush_seconds: float = Field(default=DEFAULT_IDLE_FLUSH_SECONDS, ge=0, le=30)
|
||||
|
||||
|
||||
class TerminateJobRequest(ShellctlModel):
|
||||
"""HTTP request body for `POST /v1/jobs/{job_id}/terminate`."""
|
||||
|
||||
grace_seconds: float = Field(default=DEFAULT_TERMINATE_GRACE_SECONDS, ge=0, le=300)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"TERMINAL_JOB_STATUSES",
|
||||
"DeleteJobResponse",
|
||||
"ErrorDetail",
|
||||
"ErrorResponse",
|
||||
"HealthResponse",
|
||||
"InputJobRequest",
|
||||
"JobInfo",
|
||||
"JobResult",
|
||||
"JobStatusName",
|
||||
"JobStatusView",
|
||||
"ListJobsResponse",
|
||||
"RunJobRequest",
|
||||
"ShellctlModel",
|
||||
"TerminalSize",
|
||||
"TerminateJobRequest",
|
||||
"WaitJobRequest",
|
||||
]
|
||||
7
dify-agent/src/shellctl_runtime/__init__.py
Normal file
7
dify-agent/src/shellctl_runtime/__init__.py
Normal file
@ -0,0 +1,7 @@
|
||||
"""Stdlib-only shellctl subprocess entrypoints.
|
||||
|
||||
This package is reserved for tmux hot-path helpers that must start quickly for
|
||||
every job. Keep `__init__` import-free so `shellctl-sanitize-pty` and
|
||||
`shellctl-runner-exit` do not accidentally pull in the main shellctl client or
|
||||
server stacks.
|
||||
"""
|
||||
16
dify-agent/src/shellctl_runtime/paths.py
Normal file
16
dify-agent/src/shellctl_runtime/paths.py
Normal file
@ -0,0 +1,16 @@
|
||||
"""Stdlib-only filesystem helpers shared by shellctl runtime commands."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
DEFAULT_SQLITE_BUSY_TIMEOUT_MS = 5000
|
||||
|
||||
|
||||
def db_path_from_state_dir(state_dir: Path) -> Path:
|
||||
"""Resolve the SQLite database path for a shellctl state directory."""
|
||||
|
||||
return state_dir / "shellctl.db"
|
||||
|
||||
|
||||
__all__ = ["DEFAULT_SQLITE_BUSY_TIMEOUT_MS", "db_path_from_state_dir"]
|
||||
1
dify-agent/src/shellctl_runtime/py.typed
Normal file
1
dify-agent/src/shellctl_runtime/py.typed
Normal file
@ -0,0 +1 @@
|
||||
|
||||
158
dify-agent/src/shellctl_runtime/runner_exit.py
Normal file
158
dify-agent/src/shellctl_runtime/runner_exit.py
Normal file
@ -0,0 +1,158 @@
|
||||
"""Lightweight SQLite updater for drained shellctl job exits.
|
||||
|
||||
This command runs after the tmux output pipe reaches EOF and flushes
|
||||
`output.log`. It must stay stdlib-only so exit-state materialization does not
|
||||
pay the FastAPI/SQLAlchemy import cost on every job completion. The CLI accepts
|
||||
the busy-timeout override from `ShellctlConfig` so non-default deployments keep
|
||||
the same SQLite locking behavior in and out of process.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import sqlite3
|
||||
from pathlib import Path
|
||||
|
||||
from shellctl_runtime.paths import (
|
||||
DEFAULT_SQLITE_BUSY_TIMEOUT_MS,
|
||||
db_path_from_state_dir,
|
||||
)
|
||||
|
||||
TERMINAL_STATUSES = frozenset({"exited", "terminated", "failed", "lost"})
|
||||
NONTERMINAL_STATUSES = frozenset({"created", "starting", "running"})
|
||||
|
||||
|
||||
def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
|
||||
"""Parse the tmux finalizer CLI contract."""
|
||||
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--state-dir", required=True, type=Path)
|
||||
parser.add_argument("--job-id", required=True)
|
||||
parser.add_argument("--exit-code", required=True, type=int)
|
||||
parser.add_argument("--ended-at", required=True)
|
||||
parser.add_argument(
|
||||
"--sqlite-busy-timeout-ms",
|
||||
type=int,
|
||||
default=DEFAULT_SQLITE_BUSY_TIMEOUT_MS,
|
||||
)
|
||||
return parser.parse_args(argv)
|
||||
|
||||
|
||||
def record_runner_exit(
|
||||
*,
|
||||
state_dir: Path,
|
||||
job_id: str,
|
||||
exit_code: int,
|
||||
ended_at: str,
|
||||
busy_timeout_ms: int = DEFAULT_SQLITE_BUSY_TIMEOUT_MS,
|
||||
) -> None:
|
||||
"""Persist a drained runner exit directly into the shellctl SQLite DB.
|
||||
|
||||
The update is idempotent for terminal rows: once a job reaches a terminal
|
||||
state, this helper returns successfully without rewriting the row. The
|
||||
terminal/non-terminal decision is repeated inside the `UPDATE` itself so a
|
||||
concurrent writer cannot be clobbered by a stale pre-read.
|
||||
"""
|
||||
|
||||
db_path = db_path_from_state_dir(state_dir)
|
||||
if not db_path.exists():
|
||||
raise FileNotFoundError(f"shellctl database not found: {db_path}")
|
||||
|
||||
connection = sqlite3.connect(db_path, timeout=busy_timeout_ms / 1000)
|
||||
try:
|
||||
connection.execute(f"PRAGMA busy_timeout={busy_timeout_ms}")
|
||||
row = connection.execute(
|
||||
"SELECT status FROM jobs WHERE job_id = ?",
|
||||
(job_id,),
|
||||
).fetchone()
|
||||
if row is None:
|
||||
raise LookupError(f"Unknown job id: {job_id}")
|
||||
status = str(row[0])
|
||||
if status in TERMINAL_STATUSES:
|
||||
return
|
||||
if status not in NONTERMINAL_STATUSES:
|
||||
raise ValueError(f"Unsupported shellctl job status: {status}")
|
||||
nonterminal = tuple(NONTERMINAL_STATUSES)
|
||||
cursor = connection.execute(
|
||||
"""
|
||||
UPDATE jobs
|
||||
SET status = CASE
|
||||
WHEN status IN (?, ?, ?) THEN ?
|
||||
ELSE status
|
||||
END,
|
||||
exit_code = CASE
|
||||
WHEN status IN (?, ?, ?) THEN ?
|
||||
ELSE exit_code
|
||||
END,
|
||||
ended_at = CASE
|
||||
WHEN status IN (?, ?, ?) THEN ?
|
||||
ELSE ended_at
|
||||
END,
|
||||
updated_at = CASE
|
||||
WHEN status IN (?, ?, ?) THEN ?
|
||||
ELSE updated_at
|
||||
END,
|
||||
reason = CASE
|
||||
WHEN status IN (?, ?, ?) THEN NULL
|
||||
ELSE reason
|
||||
END,
|
||||
message = CASE
|
||||
WHEN status IN (?, ?, ?) THEN NULL
|
||||
ELSE message
|
||||
END
|
||||
WHERE job_id = ?
|
||||
""",
|
||||
(
|
||||
*nonterminal,
|
||||
"exited",
|
||||
*nonterminal,
|
||||
exit_code,
|
||||
*nonterminal,
|
||||
ended_at,
|
||||
*nonterminal,
|
||||
ended_at,
|
||||
*nonterminal,
|
||||
*nonterminal,
|
||||
job_id,
|
||||
),
|
||||
)
|
||||
if cursor.rowcount == 0:
|
||||
raise LookupError(f"Unknown job id: {job_id}")
|
||||
connection.commit()
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> None:
|
||||
"""Run the standalone runner-exit command."""
|
||||
|
||||
args = parse_args(argv)
|
||||
try:
|
||||
record_runner_exit(
|
||||
state_dir=args.state_dir,
|
||||
job_id=args.job_id,
|
||||
exit_code=args.exit_code,
|
||||
ended_at=args.ended_at,
|
||||
busy_timeout_ms=args.sqlite_busy_timeout_ms,
|
||||
)
|
||||
except (
|
||||
FileNotFoundError,
|
||||
LookupError,
|
||||
OSError,
|
||||
sqlite3.DatabaseError,
|
||||
ValueError,
|
||||
) as exc:
|
||||
raise SystemExit(str(exc)) from exc
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"NONTERMINAL_STATUSES",
|
||||
"TERMINAL_STATUSES",
|
||||
"main",
|
||||
"parse_args",
|
||||
"record_runner_exit",
|
||||
]
|
||||
186
dify-agent/src/shellctl_runtime/sanitize.py
Normal file
186
dify-agent/src/shellctl_runtime/sanitize.py
Normal file
@ -0,0 +1,186 @@
|
||||
"""Lightweight PTY sanitizer used by tmux `pipe-pane`.
|
||||
|
||||
This module stays stdlib-only because shellctl starts it once per job to drain
|
||||
PTY output into `output.log`. The ready-file handshake must happen before any
|
||||
stdin reads so the server can distinguish slow startup from a stuck tmux pipe.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import codecs
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import BinaryIO
|
||||
|
||||
|
||||
class PtySanitizer:
|
||||
"""Incrementally convert PTY bytes into stable, readable UTF-8 text.
|
||||
|
||||
Responsibilities:
|
||||
- preserve UTF-8 decoder state across chunk boundaries
|
||||
- strip common ANSI control sequences without leaking partial escape state
|
||||
- normalize carriage-return progress updates into the final visible line
|
||||
|
||||
This adapter intentionally keeps only minimal terminal state. It aims for a
|
||||
practical log representation rather than a full terminal emulator.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._decoder = codecs.getincrementaldecoder("utf-8")("replace")
|
||||
self._line_buffer = ""
|
||||
self._pending_cr = False
|
||||
self._escape_state = "normal"
|
||||
|
||||
def feed(self, raw: bytes) -> str:
|
||||
"""Consume one PTY byte chunk and return newly stable text."""
|
||||
|
||||
return self._consume_text(self._decoder.decode(raw, final=False), final=False)
|
||||
|
||||
def flush(self) -> str:
|
||||
"""Flush buffered decoder/line state at end-of-stream."""
|
||||
|
||||
tail = self._consume_text(self._decoder.decode(b"", final=True), final=True)
|
||||
if self._line_buffer:
|
||||
tail += self._line_buffer
|
||||
self._line_buffer = ""
|
||||
self._pending_cr = False
|
||||
self._escape_state = "normal"
|
||||
return tail
|
||||
|
||||
def _consume_text(self, text: str, *, final: bool) -> str:
|
||||
parts: list[str] = []
|
||||
for char in text:
|
||||
self._consume_char(char, parts)
|
||||
if final and self._escape_state != "normal":
|
||||
self._escape_state = "normal"
|
||||
return "".join(parts)
|
||||
|
||||
def _consume_char(self, char: str, parts: list[str]) -> None:
|
||||
state = self._escape_state
|
||||
if state == "normal":
|
||||
if char == "\x1b":
|
||||
self._escape_state = "esc"
|
||||
return
|
||||
self._consume_visible_char(char, parts)
|
||||
return
|
||||
if state == "esc":
|
||||
if char == "[":
|
||||
self._escape_state = "csi"
|
||||
return
|
||||
if char == "]":
|
||||
self._escape_state = "osc"
|
||||
return
|
||||
self._escape_state = "normal"
|
||||
if char.isprintable() and char != "\x1b":
|
||||
self._consume_visible_char(char, parts)
|
||||
return
|
||||
if state == "csi":
|
||||
if "@" <= char <= "~":
|
||||
self._escape_state = "normal"
|
||||
return
|
||||
if state == "osc":
|
||||
if char == "\x07":
|
||||
self._escape_state = "normal"
|
||||
return
|
||||
if char == "\x1b":
|
||||
self._escape_state = "osc_esc"
|
||||
return
|
||||
if state == "osc_esc":
|
||||
self._escape_state = "normal" if char == "\\" else "osc"
|
||||
|
||||
def _consume_visible_char(self, char: str, parts: list[str]) -> None:
|
||||
if self._pending_cr:
|
||||
if char == "\n":
|
||||
parts.append(self._line_buffer)
|
||||
parts.append("\n")
|
||||
self._line_buffer = ""
|
||||
self._pending_cr = False
|
||||
return
|
||||
self._line_buffer = ""
|
||||
self._pending_cr = False
|
||||
if char == "\r":
|
||||
self._pending_cr = True
|
||||
return
|
||||
if char == "\n":
|
||||
parts.append(self._line_buffer)
|
||||
parts.append("\n")
|
||||
self._line_buffer = ""
|
||||
return
|
||||
self._line_buffer += char
|
||||
|
||||
|
||||
def sanitize_pty_output(raw: bytes) -> str:
|
||||
"""Sanitize a complete PTY byte string into readable UTF-8 text."""
|
||||
|
||||
sanitizer = PtySanitizer()
|
||||
return sanitizer.feed(raw) + sanitizer.flush()
|
||||
|
||||
|
||||
def sanitize_pty_stream(
|
||||
stdin: BinaryIO,
|
||||
stdout: BinaryIO,
|
||||
*,
|
||||
chunk_size: int = 65536,
|
||||
) -> None:
|
||||
"""Run the streaming PTY sanitizer as a Unix-style filter."""
|
||||
|
||||
sanitizer = PtySanitizer()
|
||||
while True:
|
||||
chunk = stdin.read(chunk_size)
|
||||
if not chunk:
|
||||
break
|
||||
output = sanitizer.feed(chunk)
|
||||
if output:
|
||||
stdout.write(output.encode("utf-8"))
|
||||
if hasattr(stdout, "flush"):
|
||||
stdout.flush()
|
||||
tail = sanitizer.flush()
|
||||
if tail:
|
||||
stdout.write(tail.encode("utf-8"))
|
||||
if hasattr(stdout, "flush"):
|
||||
stdout.flush()
|
||||
if hasattr(stdout, "flush"):
|
||||
stdout.flush()
|
||||
|
||||
|
||||
def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
|
||||
"""Parse the tiny CLI contract used by tmux `pipe-pane`."""
|
||||
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--ready-file", type=Path)
|
||||
return parser.parse_args(argv)
|
||||
|
||||
|
||||
def run_sanitize_pty(
|
||||
ready_file: Path | None,
|
||||
*,
|
||||
stdin: BinaryIO,
|
||||
stdout: BinaryIO,
|
||||
) -> None:
|
||||
"""Touch the ready file, then sanitize stdin into stdout."""
|
||||
|
||||
if ready_file is not None:
|
||||
ready_file.touch()
|
||||
sanitize_pty_stream(stdin, stdout)
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> None:
|
||||
"""Run the standalone PTY sanitizer module."""
|
||||
|
||||
args = parse_args(argv)
|
||||
run_sanitize_pty(args.ready_file, stdin=sys.stdin.buffer, stdout=sys.stdout.buffer)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"PtySanitizer",
|
||||
"main",
|
||||
"parse_args",
|
||||
"run_sanitize_pty",
|
||||
"sanitize_pty_output",
|
||||
"sanitize_pty_stream",
|
||||
]
|
||||
@ -459,7 +459,7 @@ class DifyLLMAdapterModelTests(unittest.IsolatedAsyncioTestCase):
|
||||
LLMResultChunk(
|
||||
model="demo-model",
|
||||
delta=LLMResultChunkDelta(
|
||||
index=1,
|
||||
index=0,
|
||||
message=AssistantPromptMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
@ -513,6 +513,72 @@ class DifyLLMAdapterModelTests(unittest.IsolatedAsyncioTestCase):
|
||||
self.assertEqual(response.parts[2].part_kind, "text")
|
||||
self.assertEqual(cast(TextPart, response.parts[2]).content, "world")
|
||||
|
||||
async def test_request_stream_assigns_fallback_ids_to_tool_calls_without_ids(self) -> None:
|
||||
def handler(_request: httpx.Request) -> httpx.Response:
|
||||
return build_stream_response(
|
||||
LLMResultChunk(
|
||||
model="demo-model",
|
||||
delta=LLMResultChunkDelta(
|
||||
index=0,
|
||||
message=AssistantPromptMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
AssistantPromptMessage.ToolCall(
|
||||
id=None,
|
||||
type="function",
|
||||
function=AssistantPromptMessage.ToolCall.ToolCallFunction(
|
||||
name="shell_run",
|
||||
arguments='{"script":"lookup find"}',
|
||||
),
|
||||
)
|
||||
],
|
||||
),
|
||||
),
|
||||
),
|
||||
LLMResultChunk(
|
||||
model="demo-model",
|
||||
delta=LLMResultChunkDelta(
|
||||
index=1,
|
||||
message=AssistantPromptMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
AssistantPromptMessage.ToolCall(
|
||||
id=None,
|
||||
type="function",
|
||||
function=AssistantPromptMessage.ToolCall.ToolCallFunction(
|
||||
name="shell_run",
|
||||
arguments='{"script":"lookup out"}',
|
||||
),
|
||||
)
|
||||
],
|
||||
),
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
async with self.mock_daemon_stream(httpx.MockTransport(handler)):
|
||||
adapter = DifyLLMAdapterModel(
|
||||
"demo-model",
|
||||
self.make_provider(),
|
||||
model_provider="openai",
|
||||
credentials={"api_key": "secret"},
|
||||
)
|
||||
|
||||
async with adapter.request_stream(
|
||||
[ModelRequest(parts=[UserPromptPart("hello")])],
|
||||
model_settings=None,
|
||||
model_request_parameters=ModelRequestParameters(),
|
||||
) as stream:
|
||||
events = [event async for event in stream]
|
||||
response = stream.get()
|
||||
|
||||
self.assertTrue(events)
|
||||
self.assertEqual([part.part_kind for part in response.parts], ["tool-call", "tool-call"])
|
||||
self.assertEqual(cast(ToolCallPart, response.parts[0]).tool_call_id, "chunk-0-tool-0")
|
||||
self.assertEqual(cast(ToolCallPart, response.parts[1]).tool_call_id, "chunk-1-tool-0")
|
||||
self.assertEqual(cast(ToolCallPart, response.parts[0]).args, '{"script":"lookup find"}')
|
||||
self.assertEqual(cast(ToolCallPart, response.parts[1]).args, '{"script":"lookup out"}')
|
||||
|
||||
async def test_request_splits_embedded_thinking_tags_into_parts(self) -> None:
|
||||
def handler(_request: httpx.Request) -> httpx.Response:
|
||||
return build_stream_response(*single_text_chunk("before<think>reasoning</think>after"))
|
||||
|
||||
@ -6,15 +6,20 @@ import asyncio
|
||||
import base64
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import cast
|
||||
|
||||
import httpx
|
||||
import httpx2 as httpx
|
||||
import pytest
|
||||
|
||||
from dify_agent.adapters.shell import shellctl
|
||||
from dify_agent.adapters.shell.config import ShellAdapterSettings
|
||||
from dify_agent.adapters.shell.factory import create_shell_provider
|
||||
from dify_agent.adapters.shell.protocols import ShellCommandResult, ShellProviderError
|
||||
from dify_agent.adapters.shell.shellctl import ShellFileTransferError, ShellctlProvider
|
||||
from dify_agent.adapters.shell.shellctl import (
|
||||
ShellctlClientProtocol,
|
||||
ShellFileTransferError,
|
||||
ShellctlProvider,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
@ -66,45 +71,73 @@ class FakeShellctlClient:
|
||||
delete_calls: list[tuple[str, bool, float | None]] = field(default_factory=list)
|
||||
closed: bool = False
|
||||
|
||||
async def run(self, script, *, cwd=None, env=None, timeout=30.0):
|
||||
async def run(
|
||||
self,
|
||||
script: str,
|
||||
*,
|
||||
cwd: str | None = None,
|
||||
env: dict[str, str] | None = None,
|
||||
timeout: float = 30.0,
|
||||
) -> _Job:
|
||||
self.run_calls.append(_RunCall(script=script, cwd=cwd, env=env, timeout=timeout))
|
||||
if self.run_handler is not None:
|
||||
return self.run_handler(script, cwd, env, timeout)
|
||||
return _Job(job_id="job", status="exited", done=True, exit_code=0)
|
||||
|
||||
async def wait(self, job_id, *, offset, timeout=30.0):
|
||||
async def wait(self, job_id: str, *, offset: int, timeout: float = 30.0) -> _Job:
|
||||
self.wait_calls.append((job_id, offset, timeout))
|
||||
if self.wait_handler is not None:
|
||||
return self.wait_handler(job_id, offset, timeout)
|
||||
return _Job(job_id=job_id, status="exited", done=True, offset=offset, exit_code=0)
|
||||
|
||||
async def input(self, job_id, text, *, offset, timeout=30.0):
|
||||
async def input(
|
||||
self,
|
||||
job_id: str,
|
||||
text: str,
|
||||
*,
|
||||
offset: int,
|
||||
timeout: float = 30.0,
|
||||
) -> _Job:
|
||||
self.input_calls.append((job_id, text, offset, timeout))
|
||||
if self.input_handler is not None:
|
||||
return self.input_handler(job_id, text, offset, timeout)
|
||||
return _Job(job_id=job_id, status="exited", done=True, offset=offset, exit_code=0)
|
||||
|
||||
async def tail(self, job_id):
|
||||
async def tail(self, job_id: str) -> _Job:
|
||||
if self.tail_handler is not None:
|
||||
return self.tail_handler(job_id)
|
||||
return _Job(job_id=job_id, status="exited", done=True, output="", exit_code=0)
|
||||
|
||||
async def terminate(self, job_id, grace_seconds=10.0):
|
||||
async def terminate(self, job_id: str, grace_seconds: float = 10.0) -> _Status:
|
||||
self.terminate_calls.append((job_id, grace_seconds))
|
||||
if self.terminate_handler is not None:
|
||||
return self.terminate_handler(job_id, grace_seconds)
|
||||
return _Status(job_id=job_id)
|
||||
|
||||
async def delete(self, job_id, *, force=False, grace_seconds=None):
|
||||
async def delete(
|
||||
self,
|
||||
job_id: str,
|
||||
*,
|
||||
force: bool = False,
|
||||
grace_seconds: float | None = None,
|
||||
) -> None:
|
||||
self.delete_calls.append((job_id, force, grace_seconds))
|
||||
return None
|
||||
|
||||
async def close(self):
|
||||
async def close(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
|
||||
def _provider(client: FakeShellctlClient) -> ShellctlProvider:
|
||||
return ShellctlProvider(entrypoint="http://shellctl", token="", client_factory=lambda: client)
|
||||
return ShellctlProvider(
|
||||
entrypoint="http://shellctl",
|
||||
token="",
|
||||
client_factory=lambda: _client_protocol(client),
|
||||
)
|
||||
|
||||
|
||||
def _client_protocol(client: FakeShellctlClient) -> ShellctlClientProtocol:
|
||||
return cast(ShellctlClientProtocol, cast(object, client))
|
||||
|
||||
|
||||
def test_factory_unknown_provider_raises() -> None:
|
||||
@ -354,7 +387,10 @@ def test_file_transfer_timeout_is_an_end_to_end_budget(monkeypatch: pytest.Monke
|
||||
client = FakeShellctlClient(run_handler=run_handler, wait_handler=wait_handler)
|
||||
|
||||
async def scenario() -> None:
|
||||
transfer = shellctl.ShellctlFileTransfer(client=client, timeout=5.0)
|
||||
transfer = shellctl.ShellctlFileTransfer(
|
||||
client=_client_protocol(client),
|
||||
timeout=5.0,
|
||||
)
|
||||
await transfer.upload(content=b"payload", remote_path="out.bin")
|
||||
|
||||
asyncio.run(scenario())
|
||||
@ -380,7 +416,10 @@ def test_file_transfer_timeout_exhaustion_raises_timeout_and_still_deletes_job(
|
||||
client = FakeShellctlClient(run_handler=run_handler)
|
||||
|
||||
async def scenario() -> None:
|
||||
transfer = shellctl.ShellctlFileTransfer(client=client, timeout=5.0)
|
||||
transfer = shellctl.ShellctlFileTransfer(
|
||||
client=_client_protocol(client),
|
||||
timeout=5.0,
|
||||
)
|
||||
with pytest.raises(ShellProviderError, match="timed out") as exc_info:
|
||||
await transfer.upload(content=b"payload", remote_path="out.bin")
|
||||
assert exc_info.value.code == "timeout"
|
||||
|
||||
@ -53,6 +53,7 @@ def test_upload_file_from_environment_requests_signed_url_and_normalizes_output(
|
||||
|
||||
def fake_request_agent_stub_file_download_sync(**kwargs):
|
||||
captured_download_request["file"] = kwargs["file"]
|
||||
captured_download_request["for_external"] = kwargs.get("for_external", True)
|
||||
return type(
|
||||
"Response",
|
||||
(),
|
||||
@ -85,6 +86,7 @@ def test_upload_file_from_environment_requests_signed_url_and_normalizes_output(
|
||||
transfer_method="tool_file",
|
||||
reference=_reference("tool-file-1"),
|
||||
)
|
||||
assert captured_download_request["for_external"] is True
|
||||
|
||||
|
||||
def test_upload_tool_file_resource_from_environment_preserves_tool_file_id(
|
||||
@ -256,6 +258,7 @@ def test_download_file_from_environment_supports_mapping_json(
|
||||
|
||||
def fake_request_download(**kwargs):
|
||||
captured["file"] = kwargs["file"]
|
||||
captured["for_external"] = kwargs["for_external"]
|
||||
return type(
|
||||
"Response",
|
||||
(),
|
||||
@ -283,10 +286,49 @@ def test_download_file_from_environment_supports_mapping_json(
|
||||
"reference": _reference("tool-file-1"),
|
||||
"url": None,
|
||||
}
|
||||
assert captured["for_external"] is False
|
||||
assert result.path == target_dir / "report.pdf"
|
||||
assert result.path.read_bytes() == b"downloaded"
|
||||
|
||||
|
||||
def test_download_file_from_environment_requests_internal_download_url(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
target_dir = tmp_path / "downloads"
|
||||
target_dir.mkdir()
|
||||
monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub")
|
||||
monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe")
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def fake_request_download(**kwargs):
|
||||
captured.update(kwargs)
|
||||
return type(
|
||||
"Response",
|
||||
(),
|
||||
{
|
||||
"filename": "report.pdf",
|
||||
"mime_type": "application/pdf",
|
||||
"size": 12,
|
||||
"download_url": "http://internal-files/report.pdf",
|
||||
},
|
||||
)()
|
||||
|
||||
monkeypatch.setattr("dify_agent.agent_stub.cli._files.request_agent_stub_file_download_sync", fake_request_download)
|
||||
monkeypatch.setattr(
|
||||
"dify_agent.agent_stub.cli._files.download_file_bytes_from_signed_url_sync",
|
||||
lambda **_kwargs: b"downloaded",
|
||||
)
|
||||
|
||||
_ = download_file_from_environment(
|
||||
transfer_method="tool_file",
|
||||
reference_or_url=_reference("tool-file-1"),
|
||||
local_dir=str(target_dir),
|
||||
)
|
||||
|
||||
assert captured["for_external"] is False
|
||||
|
||||
|
||||
def test_download_file_from_environment_requires_mapping_or_positional_pair() -> None:
|
||||
with pytest.raises(AgentStubValidationError, match="requires either --mapping or TRANSFER_METHOD REFERENCE_OR_URL"):
|
||||
_ = download_file_from_environment()
|
||||
|
||||
@ -161,6 +161,24 @@ def test_cli_plural_config_groups_expose_pull_push_and_delete(
|
||||
assert "delete" in captured.out
|
||||
|
||||
|
||||
def test_cli_config_skills_push_help_describes_skill_directory_requirements(
|
||||
capsys: pytest.CaptureFixture[str],
|
||||
) -> None:
|
||||
with pytest.raises(SystemExit) as exc_info:
|
||||
main(["config", "skills", "push", "--help"])
|
||||
|
||||
captured = capsys.readouterr()
|
||||
normalized_output = " ".join(captured.out.split())
|
||||
assert exc_info.value.code == 0
|
||||
assert "Skill directory requirements:" in captured.out
|
||||
assert "top-level SKILL.md" in captured.out
|
||||
assert "UTF-8 Markdown" in captured.out
|
||||
assert "YAML frontmatter matching this schema" in captured.out
|
||||
assert "--- name: <non-empty string> description: <string> ---" in normalized_output
|
||||
assert "Symlinked files are rejected" in captured.out
|
||||
assert ".venv and node_modules should be manually cleared before push" in normalized_output
|
||||
|
||||
|
||||
@pytest.mark.parametrize("argv", [["config", "file", "--help"], ["config", "skill", "--help"]])
|
||||
def test_cli_hidden_singular_alias_help_exposes_pull_only(
|
||||
capsys: pytest.CaptureFixture[str],
|
||||
|
||||
@ -157,7 +157,7 @@ def test_request_agent_stub_file_download_sync_posts_download_request() -> None:
|
||||
assert request.method == "POST"
|
||||
assert str(request.url) == "https://agent.example.com/agent-stub/files/download-request"
|
||||
assert json.loads(request.content) == {
|
||||
"file": {"transfer_method": "tool_file", "reference": _reference("tool-file-1")}
|
||||
"file": {"transfer_method": "tool_file", "reference": _reference("tool-file-1")},
|
||||
}
|
||||
return httpx.Response(
|
||||
200,
|
||||
@ -183,6 +183,39 @@ def test_request_agent_stub_file_download_sync_posts_download_request() -> None:
|
||||
assert response.download_url == "https://files.example.com/download"
|
||||
|
||||
|
||||
def test_request_agent_stub_file_download_sync_posts_internal_download_request() -> None:
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
assert request.method == "POST"
|
||||
assert str(request.url) == "https://agent.example.com/agent-stub/files/download-request"
|
||||
assert json.loads(request.content) == {
|
||||
"file": {"transfer_method": "tool_file", "reference": _reference("tool-file-1")},
|
||||
"for_external": False,
|
||||
}
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"filename": "report.pdf",
|
||||
"mime_type": "application/pdf",
|
||||
"size": 123,
|
||||
"download_url": "http://internal-files/report.pdf",
|
||||
},
|
||||
)
|
||||
|
||||
http_client = httpx.Client(transport=httpx.MockTransport(handler))
|
||||
try:
|
||||
response = request_agent_stub_file_download_sync(
|
||||
url="https://agent.example.com/agent-stub",
|
||||
auth_jwe="test-jwe",
|
||||
file=AgentStubFileMapping(transfer_method="tool_file", reference=_reference("tool-file-1")),
|
||||
for_external=False,
|
||||
sync_http_client=http_client,
|
||||
)
|
||||
finally:
|
||||
http_client.close()
|
||||
|
||||
assert response.download_url == "http://internal-files/report.pdf"
|
||||
|
||||
|
||||
def test_request_agent_stub_drive_manifest_sync_gets_manifest_request() -> None:
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
assert request.method == "GET"
|
||||
@ -456,12 +489,14 @@ def test_request_agent_stub_file_download_sync_dispatches_grpc_urls(monkeypatch:
|
||||
url="grpc://agent.example.com:9091",
|
||||
auth_jwe="token",
|
||||
file=file_mapping,
|
||||
for_external=False,
|
||||
)
|
||||
|
||||
assert captured == {
|
||||
"url": "grpc://agent.example.com:9091",
|
||||
"auth_jwe": "token",
|
||||
"file": file_mapping,
|
||||
"for_external": False,
|
||||
"timeout": 30.0,
|
||||
}
|
||||
assert response.download_url == "https://files.example.com/download"
|
||||
@ -613,7 +648,10 @@ def test_request_agent_stub_file_download_grpc_sync_maps_grpc_errors(monkeypatch
|
||||
grpc_module,
|
||||
"_require_conversions",
|
||||
lambda: SimpleNamespace(
|
||||
proto_file_download_request=lambda _pb2, *, file: {"file": file},
|
||||
proto_file_download_request=lambda _pb2, *, file, for_external: {
|
||||
"file": file,
|
||||
"for_external": for_external,
|
||||
},
|
||||
file_download_response_from_proto=lambda response: response,
|
||||
),
|
||||
)
|
||||
|
||||
@ -13,9 +13,14 @@ from dify_agent.agent_stub.grpc.conversions import (
|
||||
connect_response_from_proto,
|
||||
file_download_request_from_proto,
|
||||
proto_connect_request,
|
||||
proto_file_download_request,
|
||||
proto_file_download_response,
|
||||
)
|
||||
from dify_agent.agent_stub.protocol.agent_stub import AgentStubConnectResponse, AgentStubFileDownloadResponse
|
||||
from dify_agent.agent_stub.protocol.agent_stub import (
|
||||
AgentStubConnectResponse,
|
||||
AgentStubFileDownloadResponse,
|
||||
AgentStubFileMapping,
|
||||
)
|
||||
|
||||
|
||||
def _reference(record_id: str) -> str:
|
||||
@ -48,6 +53,32 @@ def test_file_download_request_from_proto_respects_optional_reference() -> None:
|
||||
|
||||
assert request.file.reference == _reference("tool-file-1")
|
||||
assert request.file.url is None
|
||||
assert request.for_external is True
|
||||
|
||||
|
||||
def test_file_download_request_from_proto_preserves_explicit_internal_audience() -> None:
|
||||
message = agent_stub_pb2.FileDownloadRequest(
|
||||
file=agent_stub_pb2.FileMapping(
|
||||
transfer_method="tool_file",
|
||||
reference=_reference("tool-file-1"),
|
||||
),
|
||||
for_external=False,
|
||||
)
|
||||
|
||||
request = file_download_request_from_proto(message)
|
||||
|
||||
assert request.for_external is False
|
||||
|
||||
|
||||
def test_proto_file_download_request_preserves_selected_audience() -> None:
|
||||
message = proto_file_download_request(
|
||||
agent_stub_pb2,
|
||||
file=AgentStubFileMapping(transfer_method="tool_file", reference=_reference("tool-file-1")),
|
||||
for_external=False,
|
||||
)
|
||||
|
||||
assert message.HasField("for_external") is True
|
||||
assert message.for_external is False
|
||||
|
||||
|
||||
def test_connect_request_from_proto_rejects_invalid_metadata_json() -> None:
|
||||
|
||||
@ -115,6 +115,41 @@ def test_dify_api_agent_stub_file_handler_injects_execution_context_for_download
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_dify_api_agent_stub_file_handler_forwards_internal_download_audience(monkeypatch) -> None:
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
assert str(request.url) == "https://api.example.com/inner/api/download/file/request"
|
||||
assert json.loads(request.content)["for_external"] is False
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"data": {
|
||||
"filename": "report.pdf",
|
||||
"mime_type": "application/pdf",
|
||||
"size": 123,
|
||||
"download_url": "http://internal-files/report.pdf",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
_patch_async_client(monkeypatch, handler)
|
||||
file_handler = DifyApiAgentStubFileRequestHandler(
|
||||
inner_api_url="https://api.example.com",
|
||||
inner_api_key="inner-secret",
|
||||
)
|
||||
|
||||
async def scenario() -> None:
|
||||
response = await file_handler.create_download_request(
|
||||
principal=_principal(),
|
||||
request=AgentStubFileDownloadRequest(
|
||||
file=AgentStubFileMapping(transfer_method="tool_file", reference=_reference("tool-file-1")),
|
||||
for_external=False,
|
||||
),
|
||||
)
|
||||
assert response.download_url == "http://internal-files/report.pdf"
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_dify_api_agent_stub_file_handler_rejects_missing_user_id() -> None:
|
||||
file_handler = DifyApiAgentStubFileRequestHandler(
|
||||
inner_api_url="https://api.example.com",
|
||||
|
||||
@ -130,6 +130,7 @@ def test_agent_stub_grpc_transport_delegates_file_download_requests() -> None:
|
||||
async def create_download_request(self, *, principal, request):
|
||||
assert principal.execution_context.user_id == "user-1"
|
||||
assert request.file.reference == _reference("tool-file-1")
|
||||
assert request.for_external is False
|
||||
return type(
|
||||
"Response",
|
||||
(),
|
||||
@ -158,7 +159,8 @@ def test_agent_stub_grpc_transport_delegates_file_download_requests() -> None:
|
||||
file=agent_stub_pb2.FileMapping(
|
||||
transfer_method="tool_file",
|
||||
reference=_reference("tool-file-1"),
|
||||
)
|
||||
),
|
||||
for_external=False,
|
||||
),
|
||||
metadata=(("authorization", f"Bearer {token}"),),
|
||||
)
|
||||
|
||||
@ -6,7 +6,7 @@ import httpx
|
||||
import pytest
|
||||
from pydantic import JsonValue
|
||||
from pydantic_ai import Tool
|
||||
from pydantic_ai.exceptions import UnexpectedModelBehavior
|
||||
from pydantic_ai.exceptions import ModelHTTPError, UnexpectedModelBehavior
|
||||
from pydantic_ai.messages import (
|
||||
ToolReturnPart,
|
||||
ModelMessage,
|
||||
@ -63,8 +63,8 @@ from dify_agent.protocol.schemas import (
|
||||
)
|
||||
from dify_agent.runtime.event_sink import InMemoryRunEventSink
|
||||
from dify_agent.runtime.compositor_factory import create_default_layer_providers
|
||||
from dify_agent.runtime.runner import AgentRunRunner, AgentRunValidationError
|
||||
from shell_session_manager.shellctl.shared import DeleteJobResponse, JobResult, JobStatusName, JobStatusView
|
||||
from dify_agent.runtime.runner import AgentRunRunner, AgentRunValidationError, _run_failed_error_payload
|
||||
from shellctl.shared import DeleteJobResponse, JobResult, JobStatusName, JobStatusView
|
||||
|
||||
|
||||
class StaticToolsTestLayer(ToolsLayer):
|
||||
@ -127,6 +127,28 @@ class FakeRunnerShellctlClient:
|
||||
return DeleteJobResponse(job_id=job_id)
|
||||
|
||||
|
||||
def test_run_failed_error_payload_preserves_plugin_rate_limit_error() -> None:
|
||||
exc = ModelHTTPError(
|
||||
429,
|
||||
"gpt-4o-mini",
|
||||
{"error_type": "InvokeRateLimitError", "message": "quota exceeded"},
|
||||
)
|
||||
|
||||
message, reason = _run_failed_error_payload(exc)
|
||||
|
||||
assert message == "quota exceeded"
|
||||
assert reason == "InvokeRateLimitError"
|
||||
|
||||
|
||||
def test_run_failed_error_payload_infers_rate_limit_reason_from_status_code() -> None:
|
||||
exc = ModelHTTPError(429, "gpt-4o-mini", {"message": "too many requests"})
|
||||
|
||||
message, reason = _run_failed_error_payload(exc)
|
||||
|
||||
assert message == "too many requests"
|
||||
assert reason == "InvokeRateLimitError"
|
||||
|
||||
|
||||
def _request(
|
||||
user: str | list[str] = "hello",
|
||||
*,
|
||||
|
||||
@ -121,8 +121,8 @@ def test_protocol_and_dify_plugin_exports_do_not_import_server_only_modules() ->
|
||||
"openai",
|
||||
"pydantic_settings",
|
||||
"redis",
|
||||
"shell_session_manager.shellctl.client",
|
||||
"shell_session_manager.shellctl.server",
|
||||
"shellctl.client",
|
||||
"shellctl.server",
|
||||
],
|
||||
imports=[
|
||||
"dify_agent.protocol",
|
||||
@ -158,7 +158,7 @@ def test_agent_stub_cli_main_import_is_client_safe() -> None:
|
||||
"jwcrypto",
|
||||
"pydantic_settings",
|
||||
"redis",
|
||||
"shell_session_manager",
|
||||
"shellctl.server",
|
||||
],
|
||||
imports=[
|
||||
"dify_agent.agent_stub.client",
|
||||
@ -238,7 +238,7 @@ def test_agent_stub_cli_help_render_does_not_load_server_modules() -> None:
|
||||
"jwcrypto",
|
||||
"pydantic_settings",
|
||||
"redis",
|
||||
"shell_session_manager",
|
||||
"shellctl.server",
|
||||
]
|
||||
script = "\n".join(
|
||||
[
|
||||
@ -269,6 +269,28 @@ def test_agent_stub_cli_help_render_does_not_load_server_modules() -> None:
|
||||
_run_python_script(script)
|
||||
|
||||
|
||||
def test_shellctl_client_imports_do_not_import_server_modules() -> None:
|
||||
_run_import_check(
|
||||
blocked_imports=[
|
||||
"aiosqlite",
|
||||
"fastapi",
|
||||
"sqlalchemy",
|
||||
"sqlmodel",
|
||||
"shellctl.server.api",
|
||||
"shellctl.server.service",
|
||||
"shellctl.server.tmux",
|
||||
"uvicorn",
|
||||
],
|
||||
imports=["shellctl", "shellctl.client", "shellctl.shared", "shellctl.cli"],
|
||||
assertions=[
|
||||
"assert hasattr(shellctl, 'ShellctlClient')",
|
||||
"assert hasattr(shellctl_client, 'ShellctlClient')",
|
||||
"assert hasattr(shellctl_shared, 'JobResult')",
|
||||
"assert hasattr(shellctl_cli, 'cli')",
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def test_server_settings_import_does_not_import_agent_stub_app() -> None:
|
||||
try:
|
||||
__import__("pydantic_settings")
|
||||
|
||||
@ -0,0 +1,32 @@
|
||||
[
|
||||
{
|
||||
"name": "utf8_passthrough",
|
||||
"chunks_hex": ["68656c6c6f20e4b896e7958cf09f99820a"],
|
||||
"expected": "hello 世界🙂\n"
|
||||
},
|
||||
{
|
||||
"name": "split_utf8_chunks",
|
||||
"chunks_hex": ["e4bda0e5a5bd20f09f", "99820a"],
|
||||
"expected": "你好 🙂\n"
|
||||
},
|
||||
{
|
||||
"name": "invalid_utf8_replacement",
|
||||
"chunks_hex": ["666f80ff6f0a"],
|
||||
"expected": "fo<66><6F>o\n"
|
||||
},
|
||||
{
|
||||
"name": "ansi_sequences_removed",
|
||||
"chunks_hex": ["1b5b33316d7265641b5b306d1b5b4b0a"],
|
||||
"expected": "red\n"
|
||||
},
|
||||
{
|
||||
"name": "carriage_return_progress",
|
||||
"chunks_hex": ["30250d3530250d313030250a"],
|
||||
"expected": "100%\n"
|
||||
},
|
||||
{
|
||||
"name": "crlf_normalized",
|
||||
"chunks_hex": ["6f6b0d0a6e6578740d0a"],
|
||||
"expected": "ok\nnext\n"
|
||||
}
|
||||
]
|
||||
688
dify-agent/tests/local/shellctl/test_shellctl_cli.py
Normal file
688
dify-agent/tests/local/shellctl/test_shellctl_cli.py
Normal file
@ -0,0 +1,688 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import ClassVar, TypeVar, cast
|
||||
|
||||
import click
|
||||
import httpx2 as httpx
|
||||
import pytest
|
||||
import typer.main
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from shellctl.client import ShellctlClientError
|
||||
from shellctl.shared.constants import DEFAULT_BASE_URL
|
||||
from shellctl.shared.schemas import (
|
||||
DeleteJobResponse,
|
||||
HealthResponse,
|
||||
JobInfo,
|
||||
JobResult,
|
||||
JobStatusName,
|
||||
JobStatusView,
|
||||
TerminalSize,
|
||||
)
|
||||
|
||||
cli_module = importlib.import_module("shellctl.cli")
|
||||
cli = cli_module.cli
|
||||
runner = CliRunner()
|
||||
ResultT = TypeVar("ResultT")
|
||||
|
||||
|
||||
def _click_command(command_name: str) -> click.Command:
|
||||
return cast(click.Group, typer.main.get_command(cli)).commands[command_name]
|
||||
|
||||
|
||||
def _command_option_names(command_name: str) -> set[str]:
|
||||
command = _click_command(command_name)
|
||||
return {
|
||||
option
|
||||
for parameter in command.params
|
||||
for option in getattr(parameter, "opts", []) # noqa: no-new-getattr click option introspection
|
||||
}
|
||||
|
||||
|
||||
def _command_option(command_name: str, option_name: str) -> click.Option:
|
||||
command = _click_command(command_name)
|
||||
for parameter in command.params:
|
||||
if isinstance(parameter, click.Option) and option_name in parameter.opts:
|
||||
return parameter
|
||||
raise AssertionError(f"{option_name} not found on {command_name}")
|
||||
|
||||
|
||||
class RecordingShellctlClient:
|
||||
init_calls: ClassVar[list[dict[str, object]]] = []
|
||||
calls: ClassVar[list[tuple[str, tuple[object, ...], dict[str, object]]]] = []
|
||||
results: ClassVar[dict[str, object]] = {}
|
||||
error: ClassVar[BaseException | None] = None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
base_url: str,
|
||||
*,
|
||||
output_limit: int,
|
||||
idle_flush_seconds: float,
|
||||
token: str | None = None,
|
||||
client: object | None = None,
|
||||
transport: object | None = None,
|
||||
) -> None:
|
||||
del client, transport
|
||||
type(self).init_calls.append(
|
||||
{
|
||||
"base_url": base_url,
|
||||
"output_limit": output_limit,
|
||||
"idle_flush_seconds": idle_flush_seconds,
|
||||
"token": token,
|
||||
}
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def reset(cls) -> None:
|
||||
cls.init_calls = []
|
||||
cls.calls = []
|
||||
cls.results = {}
|
||||
cls.error = None
|
||||
|
||||
async def __aenter__(self) -> RecordingShellctlClient:
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None:
|
||||
return None
|
||||
|
||||
def _result(self, method: str, result_type: type[ResultT]) -> ResultT:
|
||||
del result_type
|
||||
error = type(self).error
|
||||
if error is not None:
|
||||
raise error
|
||||
return cast(ResultT, type(self).results[method])
|
||||
|
||||
async def health(self) -> HealthResponse:
|
||||
type(self).calls.append(("health", (), {}))
|
||||
return self._result("health", HealthResponse)
|
||||
|
||||
async def run(
|
||||
self,
|
||||
script: str,
|
||||
*,
|
||||
cwd: str | None = None,
|
||||
env: dict[str, str] | None = None,
|
||||
timeout: float,
|
||||
terminal: TerminalSize | None = None,
|
||||
) -> JobResult:
|
||||
type(self).calls.append(
|
||||
(
|
||||
"run",
|
||||
(script,),
|
||||
{
|
||||
"cwd": cwd,
|
||||
"env": env,
|
||||
"timeout": timeout,
|
||||
"terminal": terminal,
|
||||
},
|
||||
)
|
||||
)
|
||||
return self._result("run", JobResult)
|
||||
|
||||
async def wait(self, job_id: str, *, offset: int, timeout: float) -> JobResult:
|
||||
type(self).calls.append(("wait", (job_id,), {"offset": offset, "timeout": timeout}))
|
||||
return self._result("wait", JobResult)
|
||||
|
||||
async def status(self, job_id: str) -> JobStatusView:
|
||||
type(self).calls.append(("status", (job_id,), {}))
|
||||
return self._result("status", JobStatusView)
|
||||
|
||||
async def list_jobs(
|
||||
self,
|
||||
*,
|
||||
status: JobStatusName | None = None,
|
||||
limit: int,
|
||||
) -> list[JobInfo]:
|
||||
type(self).calls.append(("list_jobs", (), {"status": status, "limit": limit}))
|
||||
return cast(list[JobInfo], self._result("list_jobs", list))
|
||||
|
||||
async def input(
|
||||
self,
|
||||
job_id: str,
|
||||
text: str,
|
||||
*,
|
||||
offset: int,
|
||||
timeout: float,
|
||||
) -> JobResult:
|
||||
type(self).calls.append(
|
||||
(
|
||||
"input",
|
||||
(job_id, text),
|
||||
{"offset": offset, "timeout": timeout},
|
||||
)
|
||||
)
|
||||
return self._result("input", JobResult)
|
||||
|
||||
async def tail(self, job_id: str) -> JobResult:
|
||||
type(self).calls.append(("tail", (job_id,), {}))
|
||||
return self._result("tail", JobResult)
|
||||
|
||||
async def terminate(self, job_id: str, grace_seconds: float) -> JobStatusView:
|
||||
type(self).calls.append(("terminate", (job_id,), {"grace_seconds": grace_seconds}))
|
||||
return self._result("terminate", JobStatusView)
|
||||
|
||||
async def delete(
|
||||
self,
|
||||
job_id: str,
|
||||
*,
|
||||
force: bool = False,
|
||||
grace_seconds: float | None = None,
|
||||
) -> DeleteJobResponse:
|
||||
type(self).calls.append(
|
||||
(
|
||||
"delete",
|
||||
(job_id,),
|
||||
{"force": force, "grace_seconds": grace_seconds},
|
||||
)
|
||||
)
|
||||
return self._result("delete", DeleteJobResponse)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def patched_client(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> type[RecordingShellctlClient]:
|
||||
RecordingShellctlClient.reset()
|
||||
monkeypatch.setattr(cli_module, "ShellctlClient", RecordingShellctlClient)
|
||||
return RecordingShellctlClient
|
||||
|
||||
|
||||
def _package_env() -> dict[str, str]:
|
||||
package_root = Path(__file__).resolve().parents[3]
|
||||
src_path = package_root / "src"
|
||||
env = dict(os.environ)
|
||||
current_pythonpath = env.get("PYTHONPATH")
|
||||
env["PYTHONPATH"] = f"{src_path}{os.pathsep}{current_pythonpath}" if current_pythonpath else str(src_path)
|
||||
return env
|
||||
|
||||
|
||||
def test_shellctl_help_lists_network_commands() -> None:
|
||||
result = runner.invoke(cli, ["--help"])
|
||||
|
||||
assert result.exit_code == 0, result.stderr
|
||||
assert "health" in result.stdout
|
||||
assert "run" in result.stdout
|
||||
assert "wait" in result.stdout
|
||||
assert "status" in result.stdout
|
||||
assert "list" in result.stdout
|
||||
assert "input" in result.stdout
|
||||
assert "tail" in result.stdout
|
||||
assert "terminate" in result.stdout
|
||||
assert "delete" in result.stdout
|
||||
assert "serve" in result.stdout
|
||||
|
||||
|
||||
def test_shellctl_run_and_serve_help_show_the_new_option_boundaries() -> None:
|
||||
run_result = runner.invoke(cli, ["run", "--help"])
|
||||
serve_result = runner.invoke(cli, ["serve", "--help"])
|
||||
|
||||
assert run_result.exit_code == 0, run_result.stderr
|
||||
assert "--base-url" in run_result.stdout
|
||||
assert "--auth-token" in run_result.stdout
|
||||
assert "--state-dir" not in run_result.stdout
|
||||
assert "--runtime-dir" not in run_result.stdout
|
||||
|
||||
assert serve_result.exit_code == 0, serve_result.stderr
|
||||
assert "--state-dir" in serve_result.stdout
|
||||
assert "--runtime-dir" in serve_result.stdout
|
||||
assert "--auth-token" in serve_result.stdout
|
||||
|
||||
|
||||
def test_shellctl_health_uses_sdk_health_and_ignores_auth_token_input(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
patched_client: type[RecordingShellctlClient],
|
||||
) -> None:
|
||||
patched_client.results["health"] = HealthResponse(status="ok")
|
||||
monkeypatch.setenv("SHELLCTL_AUTH_TOKEN", "from-env")
|
||||
|
||||
result = runner.invoke(cli, ["health", "--auth-token", "flag-token"])
|
||||
|
||||
assert result.exit_code == 0, result.stderr
|
||||
assert json.loads(result.stdout) == {"status": "ok"}
|
||||
assert patched_client.calls == [("health", (), {})]
|
||||
assert patched_client.init_calls == [
|
||||
{
|
||||
"base_url": DEFAULT_BASE_URL,
|
||||
"output_limit": 8192,
|
||||
"idle_flush_seconds": 0.5,
|
||||
"token": None,
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_shellctl_run_builds_sdk_request_and_emits_json(
|
||||
patched_client: type[RecordingShellctlClient],
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
patched_client.results["run"] = JobResult(
|
||||
job_id="job-run",
|
||||
done=False,
|
||||
status=JobStatusName.RUNNING,
|
||||
output_path="/tmp/output.log",
|
||||
output="hello\n",
|
||||
offset=6,
|
||||
truncated=False,
|
||||
)
|
||||
|
||||
result = runner.invoke(
|
||||
cli,
|
||||
[
|
||||
"run",
|
||||
"printf hello\\n",
|
||||
"--cwd",
|
||||
str(tmp_path / "workspace"),
|
||||
"--env",
|
||||
"A=1",
|
||||
"--env",
|
||||
"EMPTY=",
|
||||
"--timeout",
|
||||
"12",
|
||||
"--output-limit",
|
||||
"4096",
|
||||
"--idle-flush-seconds",
|
||||
"0.25",
|
||||
"--cols",
|
||||
"90",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.stderr
|
||||
assert json.loads(result.stdout) == {
|
||||
"job_id": "job-run",
|
||||
"done": False,
|
||||
"status": "running",
|
||||
"output_path": "/tmp/output.log",
|
||||
"output": "hello\n",
|
||||
"offset": 6,
|
||||
"truncated": False,
|
||||
}
|
||||
assert patched_client.init_calls == [
|
||||
{
|
||||
"base_url": DEFAULT_BASE_URL,
|
||||
"output_limit": 4096,
|
||||
"idle_flush_seconds": 0.25,
|
||||
"token": None,
|
||||
}
|
||||
]
|
||||
assert patched_client.calls[0][0] == "run"
|
||||
assert patched_client.calls[0][1] == ("printf hello\\n",)
|
||||
assert patched_client.calls[0][2]["cwd"] == str(tmp_path / "workspace")
|
||||
assert patched_client.calls[0][2]["env"] == {"A": "1", "EMPTY": ""}
|
||||
assert patched_client.calls[0][2]["timeout"] == 12.0
|
||||
terminal = patched_client.calls[0][2]["terminal"]
|
||||
assert isinstance(terminal, TerminalSize)
|
||||
assert terminal.cols == 90
|
||||
assert terminal.rows == 80
|
||||
|
||||
|
||||
def test_shellctl_wait_and_input_require_offset() -> None:
|
||||
assert _command_option("wait", "--offset").required is True
|
||||
assert _command_option("input", "--offset").required is True
|
||||
|
||||
|
||||
def test_shellctl_wait_and_input_map_requests(
|
||||
patched_client: type[RecordingShellctlClient],
|
||||
) -> None:
|
||||
patched_client.results["wait"] = JobResult(
|
||||
job_id="job-1",
|
||||
done=False,
|
||||
status=JobStatusName.RUNNING,
|
||||
output_path="/tmp/wait.log",
|
||||
output="chunk",
|
||||
offset=5,
|
||||
truncated=False,
|
||||
)
|
||||
|
||||
wait_result = runner.invoke(
|
||||
cli,
|
||||
[
|
||||
"wait",
|
||||
"job-1",
|
||||
"--offset",
|
||||
"3",
|
||||
"--timeout",
|
||||
"9",
|
||||
"--output-limit",
|
||||
"2048",
|
||||
"--idle-flush-seconds",
|
||||
"0.1",
|
||||
],
|
||||
)
|
||||
|
||||
assert wait_result.exit_code == 0, wait_result.stderr
|
||||
assert patched_client.calls[0] == (
|
||||
"wait",
|
||||
("job-1",),
|
||||
{"offset": 3, "timeout": 9.0},
|
||||
)
|
||||
assert patched_client.init_calls[0]["output_limit"] == 2048
|
||||
assert patched_client.init_calls[0]["idle_flush_seconds"] == 0.1
|
||||
|
||||
RecordingShellctlClient.reset()
|
||||
patched_client.results["input"] = JobResult(
|
||||
job_id="job-1",
|
||||
done=False,
|
||||
status=JobStatusName.RUNNING,
|
||||
output_path="/tmp/input.log",
|
||||
output="reply",
|
||||
offset=8,
|
||||
truncated=False,
|
||||
)
|
||||
|
||||
input_result = runner.invoke(
|
||||
cli,
|
||||
[
|
||||
"input",
|
||||
"job-1",
|
||||
"hello\n",
|
||||
"--offset",
|
||||
"5",
|
||||
"--timeout",
|
||||
"4",
|
||||
"--output-limit",
|
||||
"512",
|
||||
"--idle-flush-seconds",
|
||||
"0",
|
||||
],
|
||||
)
|
||||
|
||||
assert input_result.exit_code == 0, input_result.stderr
|
||||
assert patched_client.calls[0] == (
|
||||
"input",
|
||||
("job-1", "hello\n"),
|
||||
{"offset": 5, "timeout": 4.0},
|
||||
)
|
||||
assert patched_client.init_calls[0]["output_limit"] == 512
|
||||
assert patched_client.init_calls[0]["idle_flush_seconds"] == 0.0
|
||||
|
||||
|
||||
def test_shellctl_list_tail_status_terminate_and_delete_map_arguments(
|
||||
patched_client: type[RecordingShellctlClient],
|
||||
) -> None:
|
||||
patched_client.results["list_jobs"] = [
|
||||
JobInfo(
|
||||
job_id="job-2",
|
||||
status=JobStatusName.RUNNING,
|
||||
created_at="2026-05-21T15:30:12Z",
|
||||
)
|
||||
]
|
||||
|
||||
list_result = runner.invoke(cli, ["list", "--status", "running", "--limit", "5"])
|
||||
|
||||
assert list_result.exit_code == 0, list_result.stderr
|
||||
assert json.loads(list_result.stdout) == [
|
||||
{
|
||||
"job_id": "job-2",
|
||||
"status": "running",
|
||||
"created_at": "2026-05-21T15:30:12Z",
|
||||
}
|
||||
]
|
||||
assert patched_client.calls[0] == (
|
||||
"list_jobs",
|
||||
(),
|
||||
{"status": JobStatusName.RUNNING, "limit": 5},
|
||||
)
|
||||
|
||||
RecordingShellctlClient.reset()
|
||||
patched_client.results["tail"] = JobResult(
|
||||
job_id="job-2",
|
||||
done=False,
|
||||
status=JobStatusName.RUNNING,
|
||||
output_path="/tmp/tail.log",
|
||||
output="tail",
|
||||
offset=4,
|
||||
truncated=False,
|
||||
)
|
||||
tail_result = runner.invoke(cli, ["tail", "job-2", "--output-limit", "16"])
|
||||
assert tail_result.exit_code == 0, tail_result.stderr
|
||||
assert patched_client.calls[0] == ("tail", ("job-2",), {})
|
||||
assert patched_client.init_calls[0]["output_limit"] == 16
|
||||
|
||||
RecordingShellctlClient.reset()
|
||||
patched_client.results["status"] = JobStatusView(
|
||||
job_id="job-2",
|
||||
status=JobStatusName.RUNNING,
|
||||
done=False,
|
||||
created_at="2026-05-21T15:30:12Z",
|
||||
started_at="2026-05-21T15:30:13Z",
|
||||
offset=4,
|
||||
)
|
||||
status_result = runner.invoke(cli, ["status", "job-2"])
|
||||
assert status_result.exit_code == 0, status_result.stderr
|
||||
assert patched_client.calls[0] == ("status", ("job-2",), {})
|
||||
|
||||
RecordingShellctlClient.reset()
|
||||
patched_client.results["terminate"] = JobStatusView(
|
||||
job_id="job-2",
|
||||
status=JobStatusName.TERMINATED,
|
||||
done=True,
|
||||
created_at="2026-05-21T15:30:12Z",
|
||||
started_at="2026-05-21T15:30:13Z",
|
||||
ended_at="2026-05-21T15:30:18Z",
|
||||
offset=4,
|
||||
)
|
||||
terminate_result = runner.invoke(
|
||||
cli,
|
||||
["terminate", "job-2", "--grace-seconds", "0.25"],
|
||||
)
|
||||
assert terminate_result.exit_code == 0, terminate_result.stderr
|
||||
assert patched_client.calls[0] == (
|
||||
"terminate",
|
||||
("job-2",),
|
||||
{"grace_seconds": 0.25},
|
||||
)
|
||||
|
||||
RecordingShellctlClient.reset()
|
||||
patched_client.results["delete"] = DeleteJobResponse(job_id="job-2")
|
||||
delete_result = runner.invoke(
|
||||
cli,
|
||||
["delete", "job-2", "--force", "--grace-seconds", "0.5"],
|
||||
)
|
||||
assert delete_result.exit_code == 0, delete_result.stderr
|
||||
assert patched_client.calls[0] == (
|
||||
"delete",
|
||||
("job-2",),
|
||||
{"force": True, "grace_seconds": 0.5},
|
||||
)
|
||||
|
||||
|
||||
def test_shellctl_run_rejects_invalid_env_entry() -> None:
|
||||
result = runner.invoke(cli, ["run", "printf bad", "--env", "MISSING_EQUALS"])
|
||||
|
||||
assert result.exit_code == 2
|
||||
assert "env entries must use NAME=VALUE format" in result.stderr
|
||||
|
||||
|
||||
def test_shellctl_commands_render_sdk_errors_as_json_stderr(
|
||||
patched_client: type[RecordingShellctlClient],
|
||||
) -> None:
|
||||
patched_client.error = ShellctlClientError(
|
||||
404,
|
||||
"job_not_found",
|
||||
"Job missing is not found",
|
||||
)
|
||||
|
||||
result = runner.invoke(cli, ["status", "missing"])
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert result.stdout == ""
|
||||
assert json.loads(result.stderr) == {
|
||||
"error": {
|
||||
"code": "job_not_found",
|
||||
"message": "Job missing is not found",
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_shellctl_commands_render_transport_errors_as_json_stderr(
|
||||
patched_client: type[RecordingShellctlClient],
|
||||
) -> None:
|
||||
patched_client.error = httpx.TransportError("connection refused")
|
||||
|
||||
result = runner.invoke(cli, ["status", "missing"])
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert result.stdout == ""
|
||||
assert json.loads(result.stderr) == {
|
||||
"error": {
|
||||
"code": "connection_error",
|
||||
"message": "connection refused",
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_shellctl_commands_render_timeouts_as_json_stderr(
|
||||
patched_client: type[RecordingShellctlClient],
|
||||
) -> None:
|
||||
patched_client.error = httpx.TimeoutException("slow server")
|
||||
|
||||
result = runner.invoke(cli, ["status", "missing"])
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert result.stdout == ""
|
||||
assert json.loads(result.stderr) == {
|
||||
"error": {
|
||||
"code": "request_timeout",
|
||||
"message": "request timed out",
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_shellctl_base_url_and_auth_token_flags_override_environment(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
patched_client: type[RecordingShellctlClient],
|
||||
) -> None:
|
||||
patched_client.results["status"] = JobStatusView(
|
||||
job_id="job-2",
|
||||
status=JobStatusName.RUNNING,
|
||||
done=False,
|
||||
created_at="2026-05-21T15:30:12Z",
|
||||
started_at="2026-05-21T15:30:13Z",
|
||||
offset=4,
|
||||
)
|
||||
monkeypatch.setenv("SHELLCTL_BASE_URL", "http://from-env:9999")
|
||||
monkeypatch.setenv("SHELLCTL_AUTH_TOKEN", "from-env-token")
|
||||
|
||||
result = runner.invoke(
|
||||
cli,
|
||||
[
|
||||
"status",
|
||||
"job-2",
|
||||
"--base-url",
|
||||
"http://override:8765",
|
||||
"--auth-token",
|
||||
"flag-token",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.stderr
|
||||
assert patched_client.init_calls == [
|
||||
{
|
||||
"base_url": "http://override:8765",
|
||||
"output_limit": 8192,
|
||||
"idle_flush_seconds": 0.5,
|
||||
"token": "flag-token",
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_shellctl_cli_controller_module_is_removed() -> None:
|
||||
assert importlib.util.find_spec("shellctl.server.cli_controller") is None
|
||||
|
||||
|
||||
def test_importing_shellctl_cli_for_run_help_skips_server_stack() -> None:
|
||||
package_root = Path(__file__).resolve().parents[3]
|
||||
command = """
|
||||
import json
|
||||
import sys
|
||||
from typer.testing import CliRunner
|
||||
from shellctl.cli import cli
|
||||
|
||||
result = CliRunner().invoke(cli, ["run", "--help"])
|
||||
print(json.dumps({
|
||||
"exit_code": result.exit_code,
|
||||
"stdout": result.stdout,
|
||||
"modules": {
|
||||
"fastapi": "fastapi" in sys.modules,
|
||||
"uvicorn": "uvicorn" in sys.modules,
|
||||
"sqlalchemy": "sqlalchemy" in sys.modules,
|
||||
"sqlmodel": "sqlmodel" in sys.modules,
|
||||
"service": "shellctl.server.service" in sys.modules,
|
||||
"api": "shellctl.server.api" in sys.modules,
|
||||
"tmux": "shellctl.server.tmux" in sys.modules,
|
||||
"shared_runtime": "shellctl.shared.runtime" in sys.modules,
|
||||
"shared_output": "shellctl.shared.output" in sys.modules,
|
||||
"shared_sanitize": "shellctl.shared.sanitize" in sys.modules,
|
||||
},
|
||||
}))
|
||||
"""
|
||||
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", command],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
cwd=package_root,
|
||||
env=_package_env(),
|
||||
)
|
||||
|
||||
assert result.returncode == 0, result.stderr
|
||||
payload = json.loads(result.stdout)
|
||||
assert payload["exit_code"] == 0
|
||||
assert "--base-url" in payload["stdout"]
|
||||
assert payload["modules"] == {
|
||||
"fastapi": False,
|
||||
"uvicorn": False,
|
||||
"sqlalchemy": False,
|
||||
"sqlmodel": False,
|
||||
"service": False,
|
||||
"api": False,
|
||||
"tmux": False,
|
||||
"shared_runtime": False,
|
||||
"shared_output": False,
|
||||
"shared_sanitize": False,
|
||||
}
|
||||
|
||||
|
||||
def test_importing_server_serve_command_skips_cli_and_sdk_modules() -> None:
|
||||
package_root = Path(__file__).resolve().parents[3]
|
||||
command = """
|
||||
import json
|
||||
import sys
|
||||
from shellctl.server import serve_command
|
||||
|
||||
print(json.dumps({
|
||||
"callable": callable(serve_command),
|
||||
"modules": {
|
||||
"cli": "shellctl.cli" in sys.modules,
|
||||
"sdk": "shellctl.client.sdk" in sys.modules,
|
||||
"client": "shellctl.client" in sys.modules,
|
||||
},
|
||||
}))
|
||||
"""
|
||||
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", command],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
cwd=package_root,
|
||||
env=_package_env(),
|
||||
)
|
||||
|
||||
assert result.returncode == 0, result.stderr
|
||||
payload = json.loads(result.stdout)
|
||||
assert payload["callable"] is True
|
||||
assert payload["modules"] == {
|
||||
"cli": False,
|
||||
"sdk": False,
|
||||
"client": False,
|
||||
}
|
||||
566
dify-agent/tests/local/shellctl/test_shellctl_client.py
Normal file
566
dify-agent/tests/local/shellctl/test_shellctl_client.py
Normal file
@ -0,0 +1,566 @@
|
||||
import json
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import ClassVar, TypedDict, cast
|
||||
|
||||
import httpx2 as httpx
|
||||
import pytest
|
||||
|
||||
from shellctl.client import ShellctlClient, ShellctlClientError
|
||||
from shellctl.client import sdk as shellctl_sdk
|
||||
from shellctl.shared import (
|
||||
DEFAULT_TERMINATE_GRACE_SECONDS,
|
||||
HealthResponse,
|
||||
JobStatusName,
|
||||
)
|
||||
|
||||
|
||||
class ForcedDeleteKwargs(TypedDict, total=False):
|
||||
force: bool
|
||||
grace_seconds: float
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize(
|
||||
("call", "expected_path", "expected_json", "wait_timeout"),
|
||||
[
|
||||
(
|
||||
lambda client: client.run(
|
||||
"printf ready\\n",
|
||||
cwd="/tmp",
|
||||
env={"HELLO": "world"},
|
||||
timeout=12,
|
||||
),
|
||||
"/v1/jobs/run",
|
||||
{
|
||||
"script": "printf ready\\n",
|
||||
"cwd": "/tmp",
|
||||
"env": {"HELLO": "world"},
|
||||
"timeout": 12.0,
|
||||
"output_limit": 4096,
|
||||
"idle_flush_seconds": 0.25,
|
||||
},
|
||||
22.0,
|
||||
),
|
||||
(
|
||||
lambda client: client.wait("job-1", offset=3, timeout=7),
|
||||
"/v1/jobs/job-1/wait",
|
||||
{
|
||||
"offset": 3,
|
||||
"timeout": 7,
|
||||
"output_limit": 4096,
|
||||
"idle_flush_seconds": 0.25,
|
||||
},
|
||||
17.0,
|
||||
),
|
||||
(
|
||||
lambda client: client.input("job-1", "ls\\n", offset=5, timeout=9),
|
||||
"/v1/jobs/job-1/input",
|
||||
{
|
||||
"text": "ls\\n",
|
||||
"offset": 5,
|
||||
"timeout": 9,
|
||||
"output_limit": 4096,
|
||||
"idle_flush_seconds": 0.25,
|
||||
},
|
||||
19.0,
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_shellctl_client_blocking_calls_use_grace_read_timeout(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
call: Callable[[ShellctlClient], Awaitable[object]],
|
||||
expected_path: str,
|
||||
expected_json: dict[str, object],
|
||||
wait_timeout: float,
|
||||
) -> None:
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
captured["method"] = request.method
|
||||
captured["path"] = request.url.path
|
||||
captured["headers"] = dict(request.headers)
|
||||
captured["json"] = json.loads(request.content.decode("utf-8"))
|
||||
captured["timeout"] = request.extensions["timeout"]
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"job_id": "05211530-k7p",
|
||||
"done": False,
|
||||
"status": "running",
|
||||
"exit_code": None,
|
||||
"output_path": "/tmp/output.log",
|
||||
"output": "",
|
||||
"offset": 0,
|
||||
"truncated": False,
|
||||
},
|
||||
)
|
||||
|
||||
monkeypatch.setenv("SHELLCTL_AUTH_TOKEN", "from-env")
|
||||
transport = httpx.MockTransport(handler)
|
||||
async with ShellctlClient(
|
||||
"http://127.0.0.1:8765",
|
||||
output_limit=4096,
|
||||
idle_flush_seconds=0.25,
|
||||
transport=transport,
|
||||
) as client:
|
||||
await call(client)
|
||||
|
||||
headers = cast(dict[str, str], captured["headers"])
|
||||
assert captured["method"] == "POST"
|
||||
assert captured["path"] == expected_path
|
||||
assert headers["authorization"] == "Bearer from-env"
|
||||
assert captured["json"] == expected_json
|
||||
assert captured["timeout"] == {
|
||||
"connect": 30.0,
|
||||
"read": wait_timeout,
|
||||
"write": 30.0,
|
||||
"pool": 30.0,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize(
|
||||
("call", "expected_method", "expected_path", "response_json", "assert_result"),
|
||||
[
|
||||
(
|
||||
lambda client: client.tail("job-1"),
|
||||
"GET",
|
||||
"/v1/jobs/job-1/log/tail",
|
||||
{
|
||||
"job_id": "job-1",
|
||||
"done": False,
|
||||
"status": "running",
|
||||
"exit_code": None,
|
||||
"output_path": "/tmp/output.log",
|
||||
"output": "tail",
|
||||
"offset": 99,
|
||||
"truncated": False,
|
||||
},
|
||||
lambda result: (result.output, result.offset) == ("tail", 99),
|
||||
),
|
||||
(
|
||||
lambda client: client.status("job-1"),
|
||||
"GET",
|
||||
"/v1/jobs/job-1",
|
||||
{
|
||||
"job_id": "job-1",
|
||||
"status": "running",
|
||||
"done": False,
|
||||
"exit_code": None,
|
||||
"created_at": "2026-01-01T00:00:00Z",
|
||||
"started_at": "2026-01-01T00:00:01Z",
|
||||
"ended_at": None,
|
||||
"offset": 99,
|
||||
},
|
||||
lambda result: (result.status, result.offset) == ("running", 99),
|
||||
),
|
||||
(
|
||||
lambda client: client.delete("job-1", force=False),
|
||||
"DELETE",
|
||||
"/v1/jobs/job-1",
|
||||
{"job_id": "job-1", "deleted": True},
|
||||
lambda result: (result.job_id, result.deleted) == ("job-1", True),
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_shellctl_client_control_calls_use_default_timeout(
|
||||
call: Callable[[ShellctlClient], Awaitable[object]],
|
||||
expected_method: str,
|
||||
expected_path: str,
|
||||
response_json: dict[str, object],
|
||||
assert_result: Callable[[object], bool],
|
||||
) -> None:
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
assert request.method == expected_method
|
||||
assert request.url.path == expected_path
|
||||
assert request.extensions["timeout"] == {
|
||||
"connect": 30.0,
|
||||
"read": 30.0,
|
||||
"write": 30.0,
|
||||
"pool": 30.0,
|
||||
}
|
||||
return httpx.Response(200, json=response_json)
|
||||
|
||||
transport = httpx.MockTransport(handler)
|
||||
async with ShellctlClient(
|
||||
"http://127.0.0.1:8765",
|
||||
output_limit=1234,
|
||||
token="secret",
|
||||
transport=transport,
|
||||
) as client:
|
||||
result = await call(client)
|
||||
|
||||
assert assert_result(result)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_shellctl_client_terminate_uses_grace_read_timeout() -> None:
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
assert request.method == "POST"
|
||||
assert request.url.path == "/v1/jobs/job-1/terminate"
|
||||
assert json.loads(request.content.decode("utf-8")) == {"grace_seconds": 7.0}
|
||||
assert request.extensions["timeout"] == {
|
||||
"connect": 30.0,
|
||||
"read": 17.0,
|
||||
"write": 30.0,
|
||||
"pool": 30.0,
|
||||
}
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"job_id": "job-1",
|
||||
"status": "terminated",
|
||||
"done": True,
|
||||
"exit_code": 130,
|
||||
"created_at": "2026-01-01T00:00:00Z",
|
||||
"started_at": "2026-01-01T00:00:01Z",
|
||||
"ended_at": "2026-01-01T00:00:08Z",
|
||||
"offset": 99,
|
||||
},
|
||||
)
|
||||
|
||||
transport = httpx.MockTransport(handler)
|
||||
async with ShellctlClient(
|
||||
"http://127.0.0.1:8765",
|
||||
token="secret",
|
||||
transport=transport,
|
||||
) as client:
|
||||
result = await client.terminate("job-1", grace_seconds=7.0)
|
||||
|
||||
assert (result.status, result.exit_code) == ("terminated", 130)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize(
|
||||
("kwargs", "expected_params", "expected_read_timeout"),
|
||||
[
|
||||
(
|
||||
{"force": True, "grace_seconds": 3.5},
|
||||
{"force": "true", "grace_seconds": "3.5"},
|
||||
13.5,
|
||||
),
|
||||
(
|
||||
{"force": True},
|
||||
{"force": "true"},
|
||||
DEFAULT_TERMINATE_GRACE_SECONDS + 10.0,
|
||||
),
|
||||
(
|
||||
{"force": True, "grace_seconds": 0.0},
|
||||
{"force": "true", "grace_seconds": "0.0"},
|
||||
10.0,
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_shellctl_client_forced_delete_uses_terminate_grace_timeout(
|
||||
kwargs: ForcedDeleteKwargs,
|
||||
expected_params: dict[str, str],
|
||||
expected_read_timeout: float,
|
||||
) -> None:
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
assert request.method == "DELETE"
|
||||
assert request.url.path == "/v1/jobs/job-1"
|
||||
assert dict(request.url.params) == expected_params
|
||||
assert request.extensions["timeout"] == {
|
||||
"connect": 30.0,
|
||||
"read": expected_read_timeout,
|
||||
"write": 30.0,
|
||||
"pool": 30.0,
|
||||
}
|
||||
return httpx.Response(200, json={"job_id": "job-1", "deleted": True})
|
||||
|
||||
transport = httpx.MockTransport(handler)
|
||||
async with ShellctlClient(
|
||||
"http://127.0.0.1:8765",
|
||||
token="secret",
|
||||
transport=transport,
|
||||
) as client:
|
||||
result = await client.delete("job-1", **kwargs)
|
||||
|
||||
assert (result.job_id, result.deleted) == ("job-1", True)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_shellctl_client_closes_owned_client_on_close_and_context_exit(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
async def handler(_request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(200, json={"status": "ok"})
|
||||
|
||||
class TrackingAsyncClient(httpx.AsyncClient):
|
||||
created_clients: ClassVar[list["TrackingAsyncClient"]] = []
|
||||
|
||||
def __init__(self, *args: object, **kwargs: object) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
self.close_calls = 0
|
||||
self.__class__.created_clients.append(self)
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self.close_calls += 1
|
||||
await super().aclose()
|
||||
|
||||
monkeypatch.setattr(shellctl_sdk.httpx, "AsyncClient", TrackingAsyncClient)
|
||||
close_client = ShellctlClient(
|
||||
"http://127.0.0.1:8765",
|
||||
transport=httpx.MockTransport(handler),
|
||||
)
|
||||
direct_owned_client = TrackingAsyncClient.created_clients[-1]
|
||||
await close_client.close()
|
||||
assert direct_owned_client.close_calls == 1
|
||||
|
||||
async with ShellctlClient(
|
||||
"http://127.0.0.1:8765",
|
||||
transport=httpx.MockTransport(handler),
|
||||
) as context_client:
|
||||
context_owned_client = TrackingAsyncClient.created_clients[-1]
|
||||
await context_client.healthz()
|
||||
|
||||
assert context_owned_client.close_calls == 1
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_shellctl_client_list_jobs_uses_query_params_and_auth() -> None:
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
assert request.method == "GET"
|
||||
assert request.url.path == "/v1/jobs"
|
||||
assert request.url.params["status"] == "running"
|
||||
assert request.url.params["limit"] == "5"
|
||||
assert request.headers["authorization"] == "Bearer secret"
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"jobs": [
|
||||
{
|
||||
"job_id": "job-1",
|
||||
"status": "running",
|
||||
"created_at": "2026-05-21T15:30:12Z",
|
||||
}
|
||||
]
|
||||
},
|
||||
)
|
||||
|
||||
transport = httpx.MockTransport(handler)
|
||||
async with ShellctlClient("http://127.0.0.1:8765", token="secret", transport=transport) as client:
|
||||
result = await client.list_jobs(status="running", limit=5)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].job_id == "job-1"
|
||||
assert result[0].status == JobStatusName.RUNNING
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_shellctl_client_wait_uses_body_and_omits_auth_without_token() -> None:
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
assert request.method == "POST"
|
||||
assert request.url.path == "/v1/jobs/job-1/wait"
|
||||
assert "authorization" not in request.headers
|
||||
assert json.loads(request.content.decode("utf-8")) == {
|
||||
"offset": 3,
|
||||
"timeout": 9.0,
|
||||
"output_limit": 2048,
|
||||
"idle_flush_seconds": 0.1,
|
||||
}
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"job_id": "job-1",
|
||||
"done": False,
|
||||
"status": "running",
|
||||
"exit_code": None,
|
||||
"output_path": "/tmp/wait.log",
|
||||
"output": "chunk",
|
||||
"offset": 8,
|
||||
"truncated": False,
|
||||
},
|
||||
)
|
||||
|
||||
transport = httpx.MockTransport(handler)
|
||||
async with ShellctlClient(
|
||||
"http://127.0.0.1:8765",
|
||||
output_limit=2048,
|
||||
idle_flush_seconds=0.1,
|
||||
transport=transport,
|
||||
) as client:
|
||||
result = await client.wait("job-1", offset=3, timeout=9)
|
||||
|
||||
assert result.job_id == "job-1"
|
||||
assert result.offset == 8
|
||||
assert result.output == "chunk"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_shellctl_client_input_uses_body_and_auth() -> None:
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
assert request.method == "POST"
|
||||
assert request.url.path == "/v1/jobs/job-1/input"
|
||||
assert request.headers["authorization"] == "Bearer secret"
|
||||
assert json.loads(request.content.decode("utf-8")) == {
|
||||
"text": "ls\n",
|
||||
"offset": 5,
|
||||
"timeout": 4.0,
|
||||
"output_limit": 512,
|
||||
"idle_flush_seconds": 0.0,
|
||||
}
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"job_id": "job-1",
|
||||
"done": False,
|
||||
"status": "running",
|
||||
"exit_code": None,
|
||||
"output_path": "/tmp/input.log",
|
||||
"output": "reply",
|
||||
"offset": 10,
|
||||
"truncated": False,
|
||||
},
|
||||
)
|
||||
|
||||
transport = httpx.MockTransport(handler)
|
||||
async with ShellctlClient(
|
||||
"http://127.0.0.1:8765",
|
||||
output_limit=512,
|
||||
idle_flush_seconds=0.0,
|
||||
token="secret",
|
||||
transport=transport,
|
||||
) as client:
|
||||
result = await client.input("job-1", "ls\n", offset=5, timeout=4)
|
||||
|
||||
assert result.output == "reply"
|
||||
assert result.offset == 10
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_shellctl_client_terminate_uses_body_and_auth() -> None:
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
assert request.method == "POST"
|
||||
assert request.url.path == "/v1/jobs/job-1/terminate"
|
||||
assert request.headers["authorization"] == "Bearer secret"
|
||||
assert json.loads(request.content.decode("utf-8")) == {"grace_seconds": 0.25}
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"job_id": "job-1",
|
||||
"status": "terminated",
|
||||
"done": True,
|
||||
"created_at": "2026-05-21T15:30:12Z",
|
||||
"started_at": "2026-05-21T15:30:13Z",
|
||||
"ended_at": "2026-05-21T15:30:18Z",
|
||||
"offset": 12,
|
||||
},
|
||||
)
|
||||
|
||||
transport = httpx.MockTransport(handler)
|
||||
async with ShellctlClient("http://127.0.0.1:8765", token="secret", transport=transport) as client:
|
||||
result = await client.terminate("job-1", 0.25)
|
||||
|
||||
assert result.status == JobStatusName.TERMINATED
|
||||
assert result.done is True
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_shellctl_client_delete_uses_query_params_and_auth() -> None:
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
assert request.method == "DELETE"
|
||||
assert request.url.path == "/v1/jobs/job-1"
|
||||
assert request.url.params["force"] == "true"
|
||||
assert request.url.params["grace_seconds"] == "0.5"
|
||||
assert request.headers["authorization"] == "Bearer secret"
|
||||
return httpx.Response(200, json={"job_id": "job-1", "deleted": True})
|
||||
|
||||
transport = httpx.MockTransport(handler)
|
||||
async with ShellctlClient("http://127.0.0.1:8765", token="secret", transport=transport) as client:
|
||||
result = await client.delete("job-1", force=True, grace_seconds=0.5)
|
||||
|
||||
assert result.job_id == "job-1"
|
||||
assert result.deleted is True
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_shellctl_client_raises_structured_errors() -> None:
|
||||
def handler(_request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(
|
||||
409,
|
||||
json={"error": {"code": "job_not_running", "message": "already done"}},
|
||||
)
|
||||
|
||||
transport = httpx.MockTransport(handler)
|
||||
async with ShellctlClient("http://127.0.0.1:8765", token="secret", transport=transport) as client:
|
||||
with pytest.raises(ShellctlClientError, match="job_not_running"):
|
||||
await client.input("job-1", "ls\n", offset=0)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_shellctl_client_does_not_close_injected_client() -> None:
|
||||
async def handler(_request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(200, json={"status": "ok"})
|
||||
|
||||
class TrackingAsyncClient(httpx.AsyncClient):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(
|
||||
base_url="http://127.0.0.1:8765",
|
||||
transport=httpx.MockTransport(handler),
|
||||
)
|
||||
self.close_calls = 0
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self.close_calls += 1
|
||||
await super().aclose()
|
||||
|
||||
injected_client = TrackingAsyncClient()
|
||||
|
||||
async with ShellctlClient(
|
||||
"http://127.0.0.1:8765",
|
||||
client=injected_client,
|
||||
) as client:
|
||||
await client.healthz()
|
||||
|
||||
assert injected_client.close_calls == 0
|
||||
await injected_client.aclose()
|
||||
assert injected_client.close_calls == 1
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_shellctl_client_raises_invalid_json_errors() -> None:
|
||||
def handler(_request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(
|
||||
200,
|
||||
headers={"content-type": "application/json"},
|
||||
content=b"{",
|
||||
)
|
||||
|
||||
transport = httpx.MockTransport(handler)
|
||||
async with ShellctlClient("http://127.0.0.1:8765", transport=transport) as client:
|
||||
with pytest.raises(ShellctlClientError) as exc_info:
|
||||
await client.status("job-1")
|
||||
|
||||
assert exc_info.value.code == "invalid_json"
|
||||
assert exc_info.value.message == "{"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_shellctl_client_raises_invalid_payload_errors() -> None:
|
||||
def handler(_request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(200, json=["not", "a", "dict"])
|
||||
|
||||
transport = httpx.MockTransport(handler)
|
||||
async with ShellctlClient("http://127.0.0.1:8765", transport=transport) as client:
|
||||
with pytest.raises(ShellctlClientError) as exc_info:
|
||||
await client.healthz()
|
||||
|
||||
assert exc_info.value.code == "invalid_payload"
|
||||
assert exc_info.value.message == '["not","a","dict"]'
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_shellctl_client_health_returns_model_and_healthz_stays_compatible() -> None:
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
assert request.method == "GET"
|
||||
assert request.url.path == "/healthz"
|
||||
return httpx.Response(200, json={"status": "ok"})
|
||||
|
||||
transport = httpx.MockTransport(handler)
|
||||
async with ShellctlClient("http://127.0.0.1:8765", transport=transport) as client:
|
||||
health = await client.health()
|
||||
healthz = await client.healthz()
|
||||
|
||||
assert health == HealthResponse(status="ok")
|
||||
assert healthz == {"status": "ok"}
|
||||
1687
dify-agent/tests/local/shellctl/test_shellctl_service.py
Normal file
1687
dify-agent/tests/local/shellctl/test_shellctl_service.py
Normal file
File diff suppressed because it is too large
Load Diff
133
dify-agent/tests/local/shellctl/test_shellctl_shared.py
Normal file
133
dify-agent/tests/local/shellctl/test_shellctl_shared.py
Normal file
@ -0,0 +1,133 @@
|
||||
import json
|
||||
import re
|
||||
from datetime import UTC, datetime
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
from typing import BinaryIO, cast
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from shellctl.shared import (
|
||||
JOB_ID_ALPHABET,
|
||||
RunJobRequest,
|
||||
generate_job_id,
|
||||
read_output_window,
|
||||
tail_output_window,
|
||||
)
|
||||
from shellctl_runtime.sanitize import (
|
||||
PtySanitizer,
|
||||
sanitize_pty_output,
|
||||
sanitize_pty_stream,
|
||||
)
|
||||
|
||||
|
||||
def test_generate_job_id_matches_proposal_format() -> None:
|
||||
job_id = generate_job_id(now=datetime(2026, 5, 21, 15, 30, tzinfo=UTC))
|
||||
|
||||
assert re.fullmatch(r"05211530-[0-9abcdefghjkmnpqrstvwxyz]{3}", job_id)
|
||||
assert all(char in f"{JOB_ID_ALPHABET}-" for char in job_id[9:])
|
||||
|
||||
|
||||
def test_sanitize_pty_golden_cases() -> None:
|
||||
fixture_path = Path(__file__).with_name("golden_shellctl_sanitize.json")
|
||||
cases = json.loads(fixture_path.read_text(encoding="utf-8"))
|
||||
|
||||
for case in cases:
|
||||
chunks = [bytes.fromhex(chunk) for chunk in case["chunks_hex"]]
|
||||
sanitizer = PtySanitizer()
|
||||
streamed = "".join(sanitizer.feed(chunk) for chunk in chunks) + sanitizer.flush()
|
||||
batch = sanitize_pty_output(b"".join(chunks))
|
||||
|
||||
assert streamed == case["expected"], case["name"]
|
||||
assert batch == case["expected"], case["name"]
|
||||
|
||||
|
||||
def test_read_output_window_preserves_utf8_boundaries(tmp_path: Path) -> None:
|
||||
output_path = tmp_path / "output.log"
|
||||
output_path.write_text("A🙂B", encoding="utf-8")
|
||||
|
||||
first = read_output_window(output_path, offset=0, limit=3)
|
||||
second = read_output_window(output_path, offset=first.offset, limit=4)
|
||||
third = read_output_window(output_path, offset=second.offset, limit=8)
|
||||
|
||||
assert first.output == "A"
|
||||
assert first.offset == 1
|
||||
assert first.truncated is True
|
||||
assert second.output == "🙂"
|
||||
assert second.offset == 5
|
||||
assert second.truncated is True
|
||||
assert third.output == "B"
|
||||
assert third.offset == output_path.stat().st_size
|
||||
assert third.truncated is False
|
||||
|
||||
|
||||
def test_read_output_window_advances_past_wide_char_even_when_limit_is_smaller(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
output_path = tmp_path / "output.log"
|
||||
output_path.write_text("🙂B", encoding="utf-8")
|
||||
|
||||
first = read_output_window(output_path, offset=0, limit=1)
|
||||
second = read_output_window(output_path, offset=first.offset, limit=1)
|
||||
|
||||
assert first.output == "🙂"
|
||||
assert first.offset == len("🙂".encode())
|
||||
assert first.truncated is True
|
||||
assert second.output == "B"
|
||||
assert second.truncated is False
|
||||
|
||||
|
||||
def test_tail_output_window_skips_partial_utf8_prefix(tmp_path: Path) -> None:
|
||||
output_path = tmp_path / "output.log"
|
||||
output_path.write_text("a🙂b", encoding="utf-8")
|
||||
|
||||
tail = tail_output_window(output_path, limit=3)
|
||||
|
||||
assert tail.output == "b"
|
||||
assert tail.offset == output_path.stat().st_size
|
||||
assert tail.truncated is False
|
||||
|
||||
|
||||
def test_sanitize_pty_stream_flushes_incrementally() -> None:
|
||||
class ChunkedInput:
|
||||
def __init__(self, chunks: list[bytes]) -> None:
|
||||
self._chunks = chunks
|
||||
|
||||
def read(self, _size: int) -> bytes:
|
||||
if not self._chunks:
|
||||
return b""
|
||||
return self._chunks.pop(0)
|
||||
|
||||
class FlushCountingOutput(BytesIO):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.flush_count = 0
|
||||
|
||||
def flush(self) -> None:
|
||||
self.flush_count += 1
|
||||
super().flush()
|
||||
|
||||
stdout = FlushCountingOutput()
|
||||
sanitize_pty_stream(
|
||||
cast(BinaryIO, cast(object, ChunkedInput([b"ready\n", b"next\n"]))),
|
||||
stdout,
|
||||
chunk_size=5,
|
||||
)
|
||||
|
||||
assert stdout.getvalue() == b"ready\nnext\n"
|
||||
assert stdout.flush_count >= 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("env", "message"),
|
||||
[
|
||||
({"": "x"}, "non-empty"),
|
||||
({"A=B": "x"}, "must not contain '='"),
|
||||
({"A\x00B": "x"}, "must not contain NUL"),
|
||||
({"A": "x\x00y"}, "must not contain NUL"),
|
||||
],
|
||||
)
|
||||
def test_run_job_request_rejects_invalid_env_entries(env: dict[str, str], message: str) -> None:
|
||||
with pytest.raises(ValidationError, match=message):
|
||||
RunJobRequest(script="printf ready\n", env=env)
|
||||
@ -0,0 +1,32 @@
|
||||
[
|
||||
{
|
||||
"name": "utf8_passthrough",
|
||||
"chunks_hex": ["68656c6c6f20e4b896e7958cf09f99820a"],
|
||||
"expected": "hello 世界🙂\n"
|
||||
},
|
||||
{
|
||||
"name": "split_utf8_chunks",
|
||||
"chunks_hex": ["e4bda0e5a5bd20f09f", "99820a"],
|
||||
"expected": "你好 🙂\n"
|
||||
},
|
||||
{
|
||||
"name": "invalid_utf8_replacement",
|
||||
"chunks_hex": ["666f80ff6f0a"],
|
||||
"expected": "fo<66><6F>o\n"
|
||||
},
|
||||
{
|
||||
"name": "ansi_sequences_removed",
|
||||
"chunks_hex": ["1b5b33316d7265641b5b306d1b5b4b0a"],
|
||||
"expected": "red\n"
|
||||
},
|
||||
{
|
||||
"name": "carriage_return_progress",
|
||||
"chunks_hex": ["30250d3530250d313030250a"],
|
||||
"expected": "100%\n"
|
||||
},
|
||||
{
|
||||
"name": "crlf_normalized",
|
||||
"chunks_hex": ["6f6b0d0a6e6578740d0a"],
|
||||
"expected": "ok\nnext\n"
|
||||
}
|
||||
]
|
||||
531
dify-agent/tests/local/shellctl_runtime/test_shellctl_runtime.py
Normal file
531
dify-agent/tests/local/shellctl_runtime/test_shellctl_runtime.py
Normal file
@ -0,0 +1,531 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
import importlib.metadata
|
||||
import importlib.util
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import sqlite3
|
||||
import subprocess
|
||||
import sys
|
||||
import textwrap
|
||||
from pathlib import Path
|
||||
from types import TracebackType
|
||||
from typing import BinaryIO, TypedDict, cast
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
|
||||
from shellctl.server import ShellctlConfig, ShellctlService
|
||||
from shellctl_runtime.runner_exit import record_runner_exit
|
||||
from shellctl_runtime.sanitize import run_sanitize_pty
|
||||
|
||||
PACKAGE_ROOT = Path(__file__).resolve().parents[3]
|
||||
SRC_PATH = PACKAGE_ROOT / "src"
|
||||
HEAVY_IMPORT_ROOTS = {
|
||||
"fastapi",
|
||||
"httpx",
|
||||
"pydantic",
|
||||
"pygments",
|
||||
"rich",
|
||||
"sqlalchemy",
|
||||
"shellctl",
|
||||
"typer",
|
||||
"uvicorn",
|
||||
}
|
||||
|
||||
|
||||
class ModuleHelpStats(TypedDict):
|
||||
code: int
|
||||
elapsed: float
|
||||
deny: list[str]
|
||||
|
||||
|
||||
def _pythonpath_env() -> dict[str, str]:
|
||||
env = dict(os.environ)
|
||||
current_pythonpath = env.get("PYTHONPATH")
|
||||
env["PYTHONPATH"] = f"{SRC_PATH}{os.pathsep}{current_pythonpath}" if current_pythonpath else str(SRC_PATH)
|
||||
return env
|
||||
|
||||
|
||||
async def _initialize_database(state_dir: Path) -> None:
|
||||
service = ShellctlService(ShellctlConfig(state_dir=state_dir, runtime_dir=state_dir / "run"))
|
||||
await service.initialize_database()
|
||||
await service.shutdown()
|
||||
|
||||
|
||||
def _insert_job_row(
|
||||
*,
|
||||
state_dir: Path,
|
||||
job_id: str,
|
||||
status: str,
|
||||
exit_code: int | None = None,
|
||||
ended_at: str | None = None,
|
||||
reason: str | None = None,
|
||||
message: str | None = None,
|
||||
updated_at: str = "2026-05-21T15:30:13Z",
|
||||
) -> None:
|
||||
db_path = state_dir / "shellctl.db"
|
||||
with sqlite3.connect(db_path) as connection:
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT INTO jobs (
|
||||
job_id,
|
||||
script_path,
|
||||
output_path,
|
||||
cwd,
|
||||
terminal_cols,
|
||||
terminal_rows,
|
||||
status,
|
||||
session_name,
|
||||
pane_target,
|
||||
exit_code,
|
||||
reason,
|
||||
message,
|
||||
created_at,
|
||||
started_at,
|
||||
ended_at,
|
||||
updated_at
|
||||
)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
job_id,
|
||||
f"jobs/{job_id}/script",
|
||||
f"jobs/{job_id}/output.log",
|
||||
"/tmp",
|
||||
80,
|
||||
24,
|
||||
status,
|
||||
f"shellctl-{job_id}",
|
||||
f"shellctl-{job_id}:0.0",
|
||||
exit_code,
|
||||
reason,
|
||||
message,
|
||||
"2026-05-21T15:30:12Z",
|
||||
"2026-05-21T15:30:13Z",
|
||||
ended_at,
|
||||
updated_at,
|
||||
),
|
||||
)
|
||||
connection.commit()
|
||||
|
||||
|
||||
def _fetch_job_row(state_dir: Path, job_id: str) -> tuple[object, ...]:
|
||||
with sqlite3.connect(state_dir / "shellctl.db") as connection:
|
||||
row = connection.execute(
|
||||
"""
|
||||
SELECT status, exit_code, ended_at, updated_at, reason, message
|
||||
FROM jobs
|
||||
WHERE job_id = ?
|
||||
""",
|
||||
(job_id,),
|
||||
).fetchone()
|
||||
assert row is not None
|
||||
return row
|
||||
|
||||
|
||||
def _measure_module_help(module_name: str) -> ModuleHelpStats:
|
||||
script = textwrap.dedent(
|
||||
f"""
|
||||
import contextlib
|
||||
import io
|
||||
import json
|
||||
import runpy
|
||||
import sys
|
||||
import time
|
||||
|
||||
sys.path.insert(0, {str(SRC_PATH)!r})
|
||||
sys.argv = [{module_name!r}, "--help"]
|
||||
start = time.perf_counter()
|
||||
stdout = io.StringIO()
|
||||
stderr = io.StringIO()
|
||||
code = 0
|
||||
with contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr):
|
||||
try:
|
||||
runpy.run_module({module_name!r}, run_name="__main__")
|
||||
except SystemExit as exc:
|
||||
if isinstance(exc.code, int):
|
||||
code = exc.code
|
||||
elapsed = time.perf_counter() - start
|
||||
deny = sorted(
|
||||
name
|
||||
for name in sys.modules
|
||||
if name.split(".", 1)[0] in {sorted(HEAVY_IMPORT_ROOTS)!r}
|
||||
)
|
||||
print(json.dumps({{"code": code, "elapsed": elapsed, "deny": deny}}))
|
||||
"""
|
||||
)
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", script],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
cwd=PACKAGE_ROOT,
|
||||
)
|
||||
assert result.returncode == 0, result.stderr
|
||||
return cast(ModuleHelpStats, json.loads(result.stdout))
|
||||
|
||||
|
||||
def test_run_sanitize_pty_touches_ready_file_before_reading(tmp_path: Path) -> None:
|
||||
ready_file = tmp_path / "ready"
|
||||
output = io.BytesIO()
|
||||
|
||||
class ReadyCheckingInput:
|
||||
def __init__(self) -> None:
|
||||
self._reads = 0
|
||||
|
||||
def read(self, _size: int) -> bytes:
|
||||
self._reads += 1
|
||||
assert ready_file.exists()
|
||||
return b"before\rafter\n" if self._reads == 1 else b""
|
||||
|
||||
run_sanitize_pty(
|
||||
ready_file,
|
||||
stdin=cast(BinaryIO, cast(object, ReadyCheckingInput())),
|
||||
stdout=output,
|
||||
)
|
||||
|
||||
assert output.getvalue() == b"after\n"
|
||||
|
||||
|
||||
def test_removed_sanitizer_modules_are_not_importable() -> None:
|
||||
assert importlib.util.find_spec("shellctl.sanitize_pty") is None
|
||||
assert importlib.util.find_spec("shellctl.shared.sanitize") is None
|
||||
|
||||
|
||||
def test_removed_sanitizer_helpers_are_not_shellctl_exports() -> None:
|
||||
import shellctl as shellctl
|
||||
import shellctl.shared as shared
|
||||
|
||||
for module in (shellctl, shared):
|
||||
for name in ("PtySanitizer", "sanitize_pty_output", "sanitize_pty_stream"):
|
||||
assert not hasattr(module, name)
|
||||
|
||||
|
||||
def test_runner_exit_matches_service_state_for_running_job(tmp_path: Path) -> None:
|
||||
runtime_state_dir = tmp_path / "runtime-state"
|
||||
service_state_dir = tmp_path / "service-state"
|
||||
anyio.run(_initialize_database, runtime_state_dir)
|
||||
anyio.run(_initialize_database, service_state_dir)
|
||||
_insert_job_row(state_dir=runtime_state_dir, job_id="job-1", status="running")
|
||||
_insert_job_row(state_dir=service_state_dir, job_id="job-1", status="running")
|
||||
|
||||
record_runner_exit(
|
||||
state_dir=runtime_state_dir,
|
||||
job_id="job-1",
|
||||
exit_code=7,
|
||||
ended_at="2026-05-21T15:30:20Z",
|
||||
)
|
||||
|
||||
async def _record_with_service() -> None:
|
||||
service = ShellctlService(
|
||||
ShellctlConfig(
|
||||
state_dir=service_state_dir,
|
||||
runtime_dir=service_state_dir / "run",
|
||||
)
|
||||
)
|
||||
try:
|
||||
await service.record_runner_exit("job-1", 7, "2026-05-21T15:30:20Z")
|
||||
finally:
|
||||
await service.shutdown()
|
||||
|
||||
anyio.run(_record_with_service)
|
||||
|
||||
assert _fetch_job_row(runtime_state_dir, "job-1") == _fetch_job_row(
|
||||
service_state_dir,
|
||||
"job-1",
|
||||
)
|
||||
|
||||
|
||||
def test_runner_exit_is_idempotent_for_terminal_job(tmp_path: Path) -> None:
|
||||
state_dir = tmp_path / "state"
|
||||
anyio.run(_initialize_database, state_dir)
|
||||
_insert_job_row(
|
||||
state_dir=state_dir,
|
||||
job_id="job-2",
|
||||
status="failed",
|
||||
exit_code=99,
|
||||
ended_at="2026-05-21T15:30:18Z",
|
||||
reason="runner_failed",
|
||||
message="boom",
|
||||
updated_at="2026-05-21T15:30:18Z",
|
||||
)
|
||||
|
||||
before = _fetch_job_row(state_dir, "job-2")
|
||||
record_runner_exit(
|
||||
state_dir=state_dir,
|
||||
job_id="job-2",
|
||||
exit_code=7,
|
||||
ended_at="2026-05-21T15:30:20Z",
|
||||
)
|
||||
|
||||
assert _fetch_job_row(state_dir, "job-2") == before
|
||||
|
||||
|
||||
def test_runner_exit_uses_explicit_busy_timeout_override(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
state_dir = tmp_path / "state"
|
||||
anyio.run(_initialize_database, state_dir)
|
||||
_insert_job_row(state_dir=state_dir, job_id="job-3", status="running")
|
||||
|
||||
captured: dict[str, object] = {}
|
||||
real_connect = sqlite3.connect
|
||||
|
||||
class RecordingConnection:
|
||||
def __init__(self, connection: sqlite3.Connection) -> None:
|
||||
self._connection = connection
|
||||
|
||||
def execute(self, sql: str, parameters: tuple[object, ...] = ()) -> sqlite3.Cursor:
|
||||
if sql.startswith("PRAGMA busy_timeout="):
|
||||
captured["pragma"] = sql
|
||||
return self._connection.execute(sql, parameters)
|
||||
|
||||
def commit(self) -> None:
|
||||
self._connection.commit()
|
||||
|
||||
def close(self) -> None:
|
||||
self._connection.close()
|
||||
|
||||
def fake_connect(path: str | Path, timeout: float) -> RecordingConnection:
|
||||
captured["path"] = str(path)
|
||||
captured["timeout"] = timeout
|
||||
return RecordingConnection(real_connect(path, timeout=timeout))
|
||||
|
||||
monkeypatch.setattr("shellctl_runtime.runner_exit.sqlite3.connect", fake_connect)
|
||||
|
||||
record_runner_exit(
|
||||
state_dir=state_dir,
|
||||
job_id="job-3",
|
||||
exit_code=5,
|
||||
ended_at="2026-05-21T15:30:20Z",
|
||||
busy_timeout_ms=6789,
|
||||
)
|
||||
|
||||
assert captured["path"] == str(state_dir / "shellctl.db")
|
||||
assert captured["timeout"] == pytest.approx(6.789)
|
||||
assert captured["pragma"] == "PRAGMA busy_timeout=6789"
|
||||
|
||||
|
||||
def test_runner_exit_does_not_clobber_newer_terminal_state(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
state_dir = tmp_path / "state"
|
||||
anyio.run(_initialize_database, state_dir)
|
||||
_insert_job_row(state_dir=state_dir, job_id="job-4", status="running")
|
||||
|
||||
real_connect = sqlite3.connect
|
||||
|
||||
class RacingConnection:
|
||||
def __init__(self, connection: sqlite3.Connection) -> None:
|
||||
self._connection = connection
|
||||
self._raced = False
|
||||
|
||||
def __enter__(self) -> RacingConnection:
|
||||
self._connection.__enter__()
|
||||
return self
|
||||
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc: BaseException | None,
|
||||
tb: TracebackType | None,
|
||||
) -> bool | None:
|
||||
return self._connection.__exit__(exc_type, exc, tb)
|
||||
|
||||
def execute(self, sql: str, parameters: tuple[object, ...] = ()) -> sqlite3.Cursor:
|
||||
if sql.lstrip().startswith("UPDATE jobs") and not self._raced:
|
||||
self._raced = True
|
||||
with real_connect(state_dir / "shellctl.db") as racing_connection:
|
||||
racing_connection.execute(
|
||||
"""
|
||||
UPDATE jobs
|
||||
SET status = ?, exit_code = ?, ended_at = ?, updated_at = ?,
|
||||
reason = ?, message = ?
|
||||
WHERE job_id = ?
|
||||
""",
|
||||
(
|
||||
"failed",
|
||||
91,
|
||||
"2026-05-21T15:30:19Z",
|
||||
"2026-05-21T15:30:19Z",
|
||||
"pipe_failed",
|
||||
"concurrent winner",
|
||||
"job-4",
|
||||
),
|
||||
)
|
||||
racing_connection.commit()
|
||||
return self._connection.execute(sql, parameters)
|
||||
|
||||
def commit(self) -> None:
|
||||
self._connection.commit()
|
||||
|
||||
def close(self) -> None:
|
||||
self._connection.close()
|
||||
|
||||
def fake_connect(path: str | Path, timeout: float = 5.0, *args: object, **kwargs: object) -> RacingConnection:
|
||||
return RacingConnection(real_connect(path, timeout=timeout))
|
||||
|
||||
monkeypatch.setattr("shellctl_runtime.runner_exit.sqlite3.connect", fake_connect)
|
||||
|
||||
record_runner_exit(
|
||||
state_dir=state_dir,
|
||||
job_id="job-4",
|
||||
exit_code=7,
|
||||
ended_at="2026-05-21T15:30:20Z",
|
||||
)
|
||||
|
||||
assert _fetch_job_row(state_dir, "job-4") == (
|
||||
"failed",
|
||||
91,
|
||||
"2026-05-21T15:30:19Z",
|
||||
"2026-05-21T15:30:19Z",
|
||||
"pipe_failed",
|
||||
"concurrent winner",
|
||||
)
|
||||
|
||||
|
||||
def test_python_m_shellctl_runtime_runner_exit_updates_row(tmp_path: Path) -> None:
|
||||
state_dir = tmp_path / "state"
|
||||
anyio.run(_initialize_database, state_dir)
|
||||
_insert_job_row(state_dir=state_dir, job_id="job-cli", status="running")
|
||||
|
||||
result = subprocess.run(
|
||||
[
|
||||
sys.executable,
|
||||
"-m",
|
||||
"shellctl_runtime.runner_exit",
|
||||
"--state-dir",
|
||||
str(state_dir),
|
||||
"--job-id",
|
||||
"job-cli",
|
||||
"--exit-code",
|
||||
"7",
|
||||
"--ended-at",
|
||||
"2026-05-21T15:30:20Z",
|
||||
],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
cwd=PACKAGE_ROOT,
|
||||
env=_pythonpath_env(),
|
||||
)
|
||||
|
||||
assert result.returncode == 0, result.stderr
|
||||
assert _fetch_job_row(state_dir, "job-cli") == (
|
||||
"exited",
|
||||
7,
|
||||
"2026-05-21T15:30:20Z",
|
||||
"2026-05-21T15:30:20Z",
|
||||
None,
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def test_python_m_shellctl_runtime_runner_exit_reports_missing_job(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
state_dir = tmp_path / "state"
|
||||
anyio.run(_initialize_database, state_dir)
|
||||
|
||||
result = subprocess.run(
|
||||
[
|
||||
sys.executable,
|
||||
"-m",
|
||||
"shellctl_runtime.runner_exit",
|
||||
"--state-dir",
|
||||
str(state_dir),
|
||||
"--job-id",
|
||||
"missing-job",
|
||||
"--exit-code",
|
||||
"7",
|
||||
"--ended-at",
|
||||
"2026-05-21T15:30:20Z",
|
||||
],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
cwd=PACKAGE_ROOT,
|
||||
env=_pythonpath_env(),
|
||||
)
|
||||
|
||||
assert result.returncode != 0
|
||||
assert "Unknown job id: missing-job" in result.stderr or "Unknown job id: missing-job" in result.stdout
|
||||
|
||||
|
||||
def test_runtime_console_scripts_are_installed() -> None:
|
||||
entry_points = {
|
||||
entry_point.name: entry_point.value for entry_point in importlib.metadata.entry_points(group="console_scripts")
|
||||
}
|
||||
|
||||
assert entry_points["shellctl-sanitize-pty"] == "shellctl_runtime.sanitize:main"
|
||||
assert entry_points["shellctl-runner-exit"] == "shellctl_runtime.runner_exit:main"
|
||||
assert shutil.which("shellctl-sanitize-pty") is not None
|
||||
assert shutil.which("shellctl-runner-exit") is not None
|
||||
|
||||
|
||||
def test_runtime_modules_only_import_stdlib_or_runtime_helpers() -> None:
|
||||
forbidden_imports = {
|
||||
"anyio",
|
||||
"fastapi",
|
||||
"httpx",
|
||||
"pydantic",
|
||||
"pygments",
|
||||
"rich",
|
||||
"shellctl",
|
||||
"sqlalchemy",
|
||||
"sqlmodel",
|
||||
"typer",
|
||||
"uvicorn",
|
||||
}
|
||||
|
||||
for filename in ("sanitize.py", "runner_exit.py"):
|
||||
tree = ast.parse((SRC_PATH / "shellctl_runtime" / filename).read_text("utf-8"))
|
||||
imports: set[str] = set()
|
||||
for node in ast.walk(tree):
|
||||
if isinstance(node, ast.Import):
|
||||
imports.update(alias.name for alias in node.names)
|
||||
elif isinstance(node, ast.ImportFrom) and node.module is not None:
|
||||
imports.add(node.module)
|
||||
assert not any(
|
||||
imported == forbidden or imported.startswith(f"{forbidden}.")
|
||||
for imported in imports
|
||||
for forbidden in forbidden_imports
|
||||
), (filename, sorted(imports))
|
||||
|
||||
|
||||
def test_python_m_shellctl_runtime_sanitize_module_stays_lightweight() -> None:
|
||||
stats = _measure_module_help("shellctl_runtime.sanitize")
|
||||
|
||||
assert stats["code"] == 0
|
||||
assert stats["elapsed"] < 0.35
|
||||
assert stats["deny"] == []
|
||||
|
||||
|
||||
def test_python_m_shellctl_runtime_runner_exit_module_stays_lightweight() -> None:
|
||||
stats = _measure_module_help("shellctl_runtime.runner_exit")
|
||||
|
||||
assert stats["code"] == 0
|
||||
assert stats["elapsed"] < 0.35
|
||||
assert stats["deny"] == []
|
||||
|
||||
|
||||
def test_python_m_shellctl_runtime_sanitize_sanitizes_output(tmp_path: Path) -> None:
|
||||
ready_file = tmp_path / "ready"
|
||||
result = subprocess.run(
|
||||
[
|
||||
sys.executable,
|
||||
"-m",
|
||||
"shellctl_runtime.sanitize",
|
||||
"--ready-file",
|
||||
str(ready_file),
|
||||
],
|
||||
input=b"before\rafter\n",
|
||||
capture_output=True,
|
||||
check=False,
|
||||
cwd=PACKAGE_ROOT,
|
||||
env=_pythonpath_env(),
|
||||
)
|
||||
|
||||
assert result.returncode == 0, result.stderr.decode("utf-8", errors="replace")
|
||||
assert ready_file.exists()
|
||||
assert result.stdout.decode("utf-8") == "after\n"
|
||||
@ -7,7 +7,9 @@ from pathlib import Path
|
||||
PROJECT_ROOT = Path(__file__).resolve().parents[2]
|
||||
|
||||
CLIENT_SHARED_DTO_DEPENDENCIES = {
|
||||
"anyio>=4.12.1,<5.0.0",
|
||||
"httpx==0.28.1",
|
||||
"httpx2>=2.5.0,<3.0.0",
|
||||
"pydantic>=2.12.5,<2.13",
|
||||
"pydantic-ai-slim>=1.102.0,<2.0.0",
|
||||
"typer>=0.16.1,<0.17",
|
||||
@ -23,7 +25,13 @@ SERVER_RUNTIME_DEPENDENCIES = {
|
||||
"pydantic-ai-slim[anthropic,google,openai]>=1.85.1,<2.0.0",
|
||||
"pydantic-settings>=2.12.0,<3.0.0",
|
||||
"redis>=7.4.0,<8.0.0",
|
||||
"shell-session-manager==2.3.1",
|
||||
"uvicorn[standard]==0.46.0",
|
||||
}
|
||||
|
||||
SHELLCTL_SERVER_DEPENDENCIES = {
|
||||
"aiosqlite>=0.21.0,<1.0.0",
|
||||
"fastapi==0.136.0",
|
||||
"sqlmodel>=0.0.24,<0.1.0",
|
||||
"uvicorn[standard]==0.46.0",
|
||||
}
|
||||
|
||||
@ -54,6 +62,7 @@ def test_project_dependencies_split_client_and_server_requirements() -> None:
|
||||
assert set(project["dependencies"]) == CLIENT_SHARED_DTO_DEPENDENCIES
|
||||
assert set(project["optional-dependencies"]["grpc"]) == GRPC_RUNTIME_DEPENDENCIES
|
||||
assert set(project["optional-dependencies"]["server"]) == SERVER_RUNTIME_DEPENDENCIES
|
||||
assert set(project["optional-dependencies"]["shellctl-server"]) == SHELLCTL_SERVER_DEPENDENCIES
|
||||
assert set(pyproject["dependency-groups"]["dev"]) == DEV_DEPENDENCIES
|
||||
|
||||
|
||||
@ -62,6 +71,8 @@ def test_default_package_discovery_excludes_example_packages() -> None:
|
||||
find_config = pyproject["tool"]["setuptools"]["packages"]["find"]
|
||||
|
||||
assert find_config["where"] == ["src"]
|
||||
assert "shellctl*" in find_config["include"]
|
||||
assert "shellctl_runtime*" in find_config["include"]
|
||||
assert "agenton_examples*" not in find_config["include"]
|
||||
assert "dify_agent_examples*" not in find_config["include"]
|
||||
|
||||
@ -72,6 +83,9 @@ def test_project_declares_console_scripts() -> None:
|
||||
|
||||
assert scripts["dify-agent"] == "dify_agent.agent_stub.cli.main:main"
|
||||
assert scripts["dify-agent-stub-server"] == "dify_agent.agent_stub.server.cli:main"
|
||||
assert scripts["shellctl"] == "shellctl.cli:main"
|
||||
assert scripts["shellctl-sanitize-pty"] == "shellctl_runtime.sanitize:main"
|
||||
assert scripts["shellctl-runner-exit"] == "shellctl_runtime.runner_exit:main"
|
||||
|
||||
|
||||
def test_local_sandbox_dockerfile_installs_stub_client_and_shellctl() -> None:
|
||||
@ -80,10 +94,14 @@ def test_local_sandbox_dockerfile_installs_stub_client_and_shellctl() -> None:
|
||||
assert "ARG NODE_VERSION=22.22.1" in dockerfile
|
||||
assert "ARG PNPM_VERSION=11.9.0" in dockerfile
|
||||
assert "ARG UV_VERSION=0.8.9" in dockerfile
|
||||
assert "ARG DIFY_AGENT_TOOL_SPEC=.[grpc]" in dockerfile
|
||||
assert "ARG SHELL_SESSION_MANAGER_TOOL_SPEC=shell-session-manager" in dockerfile
|
||||
assert "ARG DIFY_AGENT_TOOL_SPEC=.[grpc,shellctl-server]" in dockerfile
|
||||
assert "SHELL_SESSION_MANAGER_TOOL_SPEC" not in dockerfile
|
||||
assert "DIFY_AGENT_STUB_DRIVE_BASE=/mnt/drive" not in dockerfile
|
||||
assert 'ENV PATH="${UV_TOOL_BIN_DIR}:${PATH}"' not in dockerfile
|
||||
assert "UV_TOOL_DIR=/opt/dify-agent-tools/envs" in dockerfile
|
||||
assert "UV_TOOL_BIN_DIR=/opt/dify-agent-tools/bin" in dockerfile
|
||||
assert 'ENV PATH="/opt/dify-agent-tools/bin:${PATH}"' in dockerfile
|
||||
assert "bash" in dockerfile
|
||||
assert "git" in dockerfile
|
||||
assert "jq" in dockerfile
|
||||
assert "openssh-client" in dockerfile
|
||||
@ -95,13 +113,8 @@ def test_local_sandbox_dockerfile_installs_stub_client_and_shellctl() -> None:
|
||||
assert "uv export --frozen --no-dev --all-extras --no-emit-project --no-hashes" in dockerfile
|
||||
assert "> /tmp/dify-agent-constraints.txt" in dockerfile
|
||||
assert '--constraints /tmp/dify-agent-constraints.txt --link-mode=copy "${DIFY_AGENT_TOOL_SPEC}"' in dockerfile
|
||||
assert (
|
||||
'--constraints /tmp/dify-agent-constraints.txt --link-mode=copy "${SHELL_SESSION_MANAGER_TOOL_SPEC}"'
|
||||
in dockerfile
|
||||
)
|
||||
assert "DIFY_AGENT_STUB_DRIVE_BASE=/mnt/drive" in dockerfile
|
||||
assert "COPY --from=tools ${UV_TOOL_DIR} ${UV_TOOL_DIR}" in dockerfile
|
||||
assert "COPY --from=tools ${UV_TOOL_BIN_DIR} ${UV_TOOL_BIN_DIR}" in dockerfile
|
||||
assert "COPY --from=tools /opt/dify-agent-tools/envs /opt/dify-agent-tools/envs" in dockerfile
|
||||
assert "COPY --from=tools /opt/dify-agent-tools/bin /opt/dify-agent-tools/bin" in dockerfile
|
||||
assert "VIRTUAL_ENV" not in dockerfile
|
||||
assert "uv sync" not in dockerfile
|
||||
assert "mkdir -p /mnt/drive" in dockerfile
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in New Issue
Block a user