mirror of
https://github.com/langgenius/dify.git
synced 2026-11-11 22:02:23 +08:00
Charles Yao
231e423431
Track backend integration, Ruff cleanup and migration reconciliation in ENG-1271. Retain required copy contracts without implementing frontend interactions.
221 lines
7.9 KiB
Python
221 lines
7.9 KiB
Python
"""A7 SDK parsing on the native loop; all outbound I/O on the origin worker."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
from dataclasses import dataclass
|
|
from math import isfinite
|
|
from time import monotonic
|
|
from typing import override
|
|
from urllib.parse import urlsplit
|
|
|
|
import httpx
|
|
import httpx2 as sdk_httpx
|
|
from dify_agent.runtime_backend.local import LocalNativeScopeClient
|
|
from shellctl.client.sdk import ShellctlClient
|
|
|
|
from core.helper import ssrf_proxy
|
|
from libs.sensitive_http_telemetry import sensitive_http_telemetry
|
|
from services.dify_builder.native_testing.sandbox.sync_bridge import (
|
|
NativeCodeLoopBridge,
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class CodeHttpRequest:
|
|
method: str
|
|
path: str
|
|
content: bytes
|
|
scope_token: str | None = None
|
|
|
|
|
|
class CodeOriginHttp:
|
|
"""Trusted fixed endpoint only; no workflow-selected URL or request headers."""
|
|
|
|
def __init__(self, *, endpoint, scope_id, generation, token, deadline, max_bytes):
|
|
parsed = urlsplit(endpoint)
|
|
if (
|
|
parsed.scheme not in ("http", "https")
|
|
or not parsed.hostname
|
|
or parsed.path not in ("", "/")
|
|
or parsed.username
|
|
or parsed.password
|
|
or parsed.query
|
|
or parsed.fragment
|
|
or re.fullmatch(r"code_[0-9a-f]{32}", scope_id) is None
|
|
or generation != 1
|
|
or not isfinite(deadline)
|
|
or not 0 < max_bytes <= 16 * 1024 * 1024
|
|
or not token
|
|
):
|
|
raise ValueError("native_code_endpoint_invalid")
|
|
self.endpoint = endpoint.rstrip("/")
|
|
self.scope_id, self.generation, self.token = scope_id, generation, token
|
|
self.deadline, self.max_bytes = deadline, max_bytes
|
|
|
|
def request(self, request: CodeHttpRequest):
|
|
prefix = f"/v1/test-scopes/{self.scope_id}"
|
|
scope = (request.method, request.path) in (
|
|
("POST", "/v1/test-scopes"),
|
|
("GET", prefix),
|
|
("DELETE", prefix),
|
|
)
|
|
job = (
|
|
(request.method == "POST" and request.path == prefix + "/jobs/run")
|
|
or (
|
|
re.fullmatch(
|
|
re.escape(prefix) + r"/jobs/[A-Za-z0-9][A-Za-z0-9_-]{0,127}/(?:input|wait)",
|
|
request.path,
|
|
)
|
|
and request.method == "POST"
|
|
)
|
|
or (
|
|
re.fullmatch(
|
|
re.escape(prefix) + r"/jobs/[A-Za-z0-9][A-Za-z0-9_-]{0,127}",
|
|
request.path,
|
|
)
|
|
and request.method == "GET"
|
|
)
|
|
)
|
|
remaining = self.deadline - monotonic()
|
|
if not (scope or job) or remaining <= 0 or len(request.content) > self.max_bytes:
|
|
raise ValueError("native_code_http_refused")
|
|
if job and not request.scope_token:
|
|
raise ValueError("native_code_scope_token_required")
|
|
headers = {
|
|
"Authorization": f"Bearer {self.token}",
|
|
"X-Scope-Generation": str(self.generation),
|
|
"Accept-Encoding": "identity",
|
|
"Content-Type": "application/json",
|
|
}
|
|
if job:
|
|
assert request.scope_token is not None
|
|
headers["X-Scope-Token"] = request.scope_token
|
|
with sensitive_http_telemetry():
|
|
response = ssrf_proxy.make_request(
|
|
request.method,
|
|
self.endpoint + request.path,
|
|
max_retries=0,
|
|
sensitive_request=True,
|
|
stream_response=True,
|
|
follow_redirects=False,
|
|
headers=headers,
|
|
content=request.content,
|
|
timeout=httpx.Timeout(remaining),
|
|
)
|
|
try:
|
|
if response.headers.get("content-encoding", "identity").lower() not in (
|
|
"",
|
|
"identity",
|
|
):
|
|
raise ValueError("native_code_encoding_refused")
|
|
content = bytearray()
|
|
for chunk in response.iter_bytes(chunk_size=4096):
|
|
if monotonic() >= self.deadline or len(content) + len(chunk) > self.max_bytes:
|
|
raise ValueError("native_code_response_unavailable")
|
|
content.extend(chunk)
|
|
if monotonic() >= self.deadline:
|
|
raise ValueError("native_code_response_expired")
|
|
return response.status_code, bytes(content)
|
|
finally:
|
|
response.close()
|
|
|
|
|
|
class CodeMailboxHttpTransport(sdk_httpx.AsyncBaseTransport):
|
|
def __init__(self, *, bridge, endpoint, scope_token):
|
|
self.bridge, self.endpoint, self.scope_token = (
|
|
bridge,
|
|
endpoint.rstrip("/"),
|
|
scope_token,
|
|
)
|
|
self.closed = False
|
|
|
|
@override
|
|
async def handle_async_request(self, request):
|
|
if self.closed or str(request.url).split("/v1/", 1)[0] != self.endpoint or request.url.query:
|
|
raise ValueError("native_code_http_origin_mismatch")
|
|
content = await request.aread()
|
|
value = CodeHttpRequest(request.method, request.url.path, content, self.scope_token)
|
|
status, body = await self.bridge.request("http", value)
|
|
return sdk_httpx.Response(status, content=body, request=request)
|
|
|
|
@override
|
|
async def aclose(self):
|
|
self.closed = True
|
|
|
|
|
|
class _ObservedJobClient(ShellctlClient):
|
|
def __init__(self, *, bridge, **kwargs):
|
|
self._bridge = bridge
|
|
self.closed = False
|
|
super().__init__(**kwargs)
|
|
|
|
@override
|
|
async def input(self, job_id, text, *, offset, timeout=15.0):
|
|
result = await super().input(job_id, text, offset=offset, timeout=timeout)
|
|
self._observe(result, job_id, offset)
|
|
return result
|
|
|
|
@override
|
|
async def wait(self, job_id, *, offset, timeout=15.0):
|
|
result = await super().wait(job_id, offset=offset, timeout=timeout)
|
|
self._observe(result, job_id, offset)
|
|
return result
|
|
|
|
def _observe(self, result, job_id, offset):
|
|
if (
|
|
result.job_id != job_id
|
|
or result.offset != offset + len(result.output.encode("utf-8"))
|
|
or result.offset > self.output_limit
|
|
):
|
|
raise ValueError("native_code_output_invalid")
|
|
self._bridge.call("offset", (job_id, offset, result.offset))
|
|
|
|
@override
|
|
async def close(self):
|
|
await super().close()
|
|
self.closed = True
|
|
|
|
|
|
class NativeCodeScopeClient(LocalNativeScopeClient):
|
|
def __init__(self, *, endpoint, auth_token, bridge: NativeCodeLoopBridge, output_bytes):
|
|
super().__init__(endpoint, auth_token)
|
|
self.bridge = bridge
|
|
self.output_bytes = output_bytes
|
|
self.clients: list[_ObservedJobClient] = []
|
|
|
|
@override
|
|
async def _request(self, method, path, *, generation=None, spec=None):
|
|
content = json.dumps(spec.model_dump(mode="json")).encode() if spec else b""
|
|
status, body = await self.bridge.request(
|
|
"create" if method == "POST" and path == "/v1/test-scopes" else "http",
|
|
CodeHttpRequest(method, path, content),
|
|
settlement=method == "DELETE",
|
|
)
|
|
if status not in (200, 201):
|
|
raise ValueError("native_code_scope_unavailable")
|
|
result = json.loads(body)
|
|
if not isinstance(result, dict):
|
|
raise TypeError("native_code_scope_response_invalid")
|
|
return result
|
|
|
|
@override
|
|
def job_client(self, handle):
|
|
transport = CodeMailboxHttpTransport(
|
|
bridge=self.bridge,
|
|
endpoint=self.endpoint,
|
|
scope_token=handle.token.get_secret_value(),
|
|
)
|
|
client = _ObservedJobClient(
|
|
bridge=self.bridge,
|
|
base_url=self.endpoint,
|
|
token=self.auth_token,
|
|
native_scope=handle,
|
|
transport=transport,
|
|
output_limit=self.output_bytes,
|
|
request_timeout_grace_seconds=0,
|
|
)
|
|
self.clients.append(client)
|
|
return client
|