From b5e4642f27f1c8cb1a28d919e7c1ff1fc454bcfc Mon Sep 17 00:00:00 2001 From: WH-2099 Date: Wed, 23 Sep 2026 09:54:32 +0000 Subject: [PATCH 1/2] fix(invocation): preserve JSON in backwards requests --- .../json_transport_test.go | 230 ++++++++++++++++++ .../io_tunnel/backwards_invocation/request.go | 7 +- .../io_tunnel/backwards_invocation/task.go | 94 ++++--- .../transaction/json_transport_test.go | 96 ++++++++ .../transaction/serverless_handler.go | 8 +- .../transaction/serverless_handler_test.go | 1 + pkg/entities/model_entities/llm.go | 8 +- pkg/entities/model_entities/llm_test.go | 21 ++ 8 files changed, 406 insertions(+), 59 deletions(-) create mode 100644 internal/core/io_tunnel/backwards_invocation/json_transport_test.go create mode 100644 internal/core/io_tunnel/backwards_invocation/transaction/json_transport_test.go diff --git a/internal/core/io_tunnel/backwards_invocation/json_transport_test.go b/internal/core/io_tunnel/backwards_invocation/json_transport_test.go new file mode 100644 index 000000000..ee09c5c2e --- /dev/null +++ b/internal/core/io_tunnel/backwards_invocation/json_transport_test.go @@ -0,0 +1,230 @@ +package backwards_invocation + +import ( + "bytes" + "encoding/json" + "fmt" + "testing" + "time" + + "github.com/langgenius/dify-plugin-daemon/internal/core/dify_invocation" + "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel/access_types" + "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" + "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/routine" + "github.com/stretchr/testify/require" +) + +type invocationJSONWriter struct { + session *session_manager.Session + frames [][]byte + done chan struct{} +} + +func (w *invocationJSONWriter) Write(event session_manager.PLUGIN_IN_STREAM_EVENT, data any) error { + w.frames = append(w.frames, w.session.Message(event, data)) + return nil +} + +func (w *invocationJSONWriter) Done() { close(w.done) } + +type invocationJSONEvent struct { + Event RequestEvent `json:"event"` + Message string `json:"message"` + Data json.RawMessage `json:"data"` +} + +func invokeJSON(t *testing.T, wire string, from access_types.PluginAccessType, allowed bool) ([]invocationJSONEvent, error) { + t.Helper() + routine.InitPool(4) + declaration := &plugin_entities.PluginDeclaration{ + PluginDeclarationWithoutAdvancedFields: plugin_entities.PluginDeclarationWithoutAdvancedFields{ + Resource: plugin_entities.PluginResourceRequirement{ + Permission: &plugin_entities.PluginPermissionRequirement{ + Tool: &plugin_entities.PluginPermissionToolRequirement{Enabled: allowed}, + Model: &plugin_entities.PluginPermissionModelRequirement{ + Enabled: allowed, LLM: allowed, TextEmbedding: allowed, Rerank: allowed, + TTS: allowed, Speech2text: allowed, Moderation: allowed, + }, + Node: &plugin_entities.PluginPermissionNodeRequirement{Enabled: allowed}, + App: &plugin_entities.PluginPermissionAppRequirement{Enabled: allowed}, + Storage: &plugin_entities.PluginPermissionStorageRequirement{Enabled: allowed, Size: 100}, + }, + }, + }, + } + session := &session_manager.Session{ID: "session", TenantID: `trusted"tenant`, UserID: "trusted-user"} + writer := &invocationJSONWriter{session: session, done: make(chan struct{})} + err := InvokeDify(declaration, from, session, writer, []byte(wire)) + if err != nil { + require.Empty(t, writer.frames) + return nil, err + } + select { + case <-writer.done: + case <-time.After(5 * time.Second): + t.Fatal("invocation did not send an end frame") + } + var events []invocationJSONEvent + for _, frame := range writer.frames { + var envelope struct { + Event session_manager.PLUGIN_IN_STREAM_EVENT `json:"event"` + Data invocationJSONEvent `json:"data"` + } + require.NoError(t, json.Unmarshal(frame, &envelope)) + require.Equal(t, session_manager.PLUGIN_IN_STREAM_EVENT_RESPONSE, envelope.Event) + events = append(events, envelope.Data) + } + require.NotEmpty(t, events) + require.Equal(t, REQUEST_EVENT_END, events[len(events)-1].Event) + return events, nil +} + +func echoDispatch[T any](handle *BackwardsInvocation) { + genericDispatchTask(handle, func(handle *BackwardsInvocation, request *T) { + handle.WriteResponse("struct", request) + }) +} + +func replaceDispatch(t *testing.T, kind dify_invocation.InvokeType, dispatch func(*BackwardsInvocation)) { + t.Helper() + previous := dispatchMapping[kind] + dispatchMapping[kind] = dispatch + t.Cleanup(func() { dispatchMapping[kind] = previous }) +} + +func invokeJSONRequest(t *testing.T, kind dify_invocation.InvokeType, request string) json.RawMessage { + t.Helper() + wire := fmt.Sprintf(`{"type":%q,"backwards_request_id":"request","request":%s}`, kind, request) + events, err := invokeJSON(t, wire, access_types.PLUGIN_ACCESS_TYPE_ENDPOINT, true) + require.NoError(t, err) + require.Len(t, events, 2) + require.Equal(t, REQUEST_EVENT_RESPONSE, events[0].Event, events[0].Message) + return events[0].Data +} + +func TestInvokeDifyPreservesOpaqueJSON(t *testing.T) { + replaceDispatch(t, dify_invocation.INVOKE_TYPE_LLM, echoDispatch[dify_invocation.InvokeLLMRequest]) + values := []string{ + `null`, `{}`, `[]`, `""`, `"signature"`, `false`, `true`, `0`, `9007199254740993`, + `[49]`, `[{"id":9007199254740993}]`, + `{"anthropic_content":[{"type":"thinking","thinking":"","signature":"original"},{"type":"tool_use","input":{"id":9007199254740993}}]}`, + `{"responses_output":[{"type":"reasoning","encrypted_content":"original","summary":[]}]}`, + } + for _, stream := range []bool{false, true} { + for _, opaque := range values { + t.Run(fmt.Sprintf("stream=%t/%s", stream, opaque), func(t *testing.T) { + payload := fmt.Sprintf(`{"provider":"p","model":"m","mode":"chat","stream":%t,"prompt_messages":[{"role":"assistant","content":null,"opaque_body":%s}]}`, stream, opaque) + response := invokeJSONRequest(t, dify_invocation.INVOKE_TYPE_LLM, payload) + var request struct { + Stream bool `json:"stream"` + PromptMessages []struct { + Content json.RawMessage `json:"content"` + Opaque json.RawMessage `json:"opaque_body"` + } `json:"prompt_messages"` + } + require.NoError(t, json.Unmarshal(response, &request)) + require.Equal(t, stream, request.Stream) + require.Equal(t, "null", string(request.PromptMessages[0].Content)) + // Compare JSON bytes without decoding numbers to float64. + var compact bytes.Buffer + require.NoError(t, json.Compact(&compact, []byte(opaque))) + require.Equal(t, compact.String(), string(request.PromptMessages[0].Opaque)) + }) + } + } +} + +func TestInvokeDifyPreservesParametersAndTrustedIdentity(t *testing.T) { + replaceDispatch(t, dify_invocation.INVOKE_TYPE_LLM, echoDispatch[dify_invocation.InvokeLLMRequest]) + payload := `{"provider":"p","model":"m","mode":"chat","tenant_id":"spoof","TENANT_ID":"spoof","user_id":"spoof","uſer_id":"spoof","TYPE":"tool","type":"tool","completion_params":{"seed":9007199254740993},"tools":[{"name":"f","parameters":{"enum":[9007199254740993]}}],"prompt_messages":[{"role":"user","content":[{"type":"image","url":"https://example.com/image","opaque_body":{"id":9007199254740993}}]},{"role":"assistant","tool_calls":[{"id":"c","type":"function","function":{"name":"f","arguments":"{\"id\":9007199254740993}"}}]}]}` + response := invokeJSONRequest(t, dify_invocation.INVOKE_TYPE_LLM, payload) + var request struct { + TenantID string `json:"tenant_id"` + UserID string `json:"user_id"` + Type string `json:"type"` + } + require.NoError(t, json.Unmarshal(response, &request)) + require.Equal(t, `trusted"tenant`, request.TenantID) + require.Equal(t, "trusted-user", request.UserID) + require.Equal(t, "llm", request.Type) + require.Contains(t, string(response), `"seed":9007199254740993`) + require.Contains(t, string(response), `"enum":[9007199254740993]`) + require.Contains(t, string(response), `"opaque_body":{"id":9007199254740993}`) + require.Contains(t, string(response), `"arguments":"{\"id\":9007199254740993}"`) + require.Contains(t, string(response), `"role":"assistant","content":null`) +} + +func TestInvokeDifyAllDispatchTypes(t *testing.T) { + cases := []struct { + kind dify_invocation.InvokeType + payload string + dispatch func(*BackwardsInvocation) + }{ + {"llm", `{"provider":"p","model":"m","mode":"chat","prompt_messages":[{"role":"user","content":"hi"}]}`, echoDispatch[dify_invocation.InvokeLLMRequest]}, + {"llm_structured_output", `{"provider":"p","model":"m","mode":"chat","prompt_messages":[{"role":"user","content":"hi"}],"structured_output_schema":{"type":"object"}}`, echoDispatch[dify_invocation.InvokeLLMWithStructuredOutputRequest]}, + {"text_embedding", `{"provider":"p","model":"m","texts":["hi"],"input_type":"query"}`, echoDispatch[dify_invocation.InvokeTextEmbeddingRequest]}, + {"multimodal_embedding", `{"provider":"p","model":"m","documents":[{"content":"hi","content_type":"text"}],"input_type":"query"}`, echoDispatch[dify_invocation.InvokeMultimodalEmbeddingRequest]}, + {"rerank", `{"provider":"p","model":"m","query":"hi","docs":["doc"],"top_n":1,"score_threshold":0.3}`, echoDispatch[dify_invocation.InvokeRerankRequest]}, + {"multimodal_rerank", `{"provider":"p","model":"m","query":{"content":"hi","content_type":"text"},"docs":[{"content":"hi","content_type":"text"}],"top_n":1}`, echoDispatch[dify_invocation.InvokeMultimodalRerankRequest]}, + {"tts", `{"provider":"p","model":"m","content_text":"hi","voice":"test"}`, echoDispatch[dify_invocation.InvokeTTSRequest]}, + {"speech2text", `{"provider":"p","model":"m","file":"abcd"}`, echoDispatch[dify_invocation.InvokeSpeech2TextRequest]}, + {"moderation", `{"provider":"p","model":"m","text":"hi"}`, echoDispatch[dify_invocation.InvokeModerationRequest]}, + {"tool", `{"provider":"p","tool_type":"builtin","tool":"test","tool_parameters":{"id":9007199254740993}}`, echoDispatch[dify_invocation.InvokeToolRequest]}, + {"app", `{"app_id":"a","inputs":{"id":9007199254740993},"response_mode":"streaming","query":"hi"}`, echoDispatch[dify_invocation.InvokeAppRequest]}, + {"node_parameter_extractor", `{"parameters":[{"name":"x","type":"number"}],"model":{"provider":"p","name":"m","mode":"chat"},"query":"hi"}`, echoDispatch[dify_invocation.InvokeParameterExtractorRequest]}, + {"node_question_classifier", `{"classes":[{"id":"x","name":"test"}],"model":{"provider":"p","name":"m","mode":"chat"},"query":"hi"}`, echoDispatch[dify_invocation.InvokeQuestionClassifierRequest]}, + {"storage", `{"opt":"set","key":"key","value":"abcd"}`, echoDispatch[dify_invocation.InvokeStorageRequest]}, + {"system_summary", `{"text":"hi"}`, echoDispatch[dify_invocation.InvokeSummaryRequest]}, + {"upload_file", `{"filename":"hi","mimetype":"text/plain"}`, echoDispatch[dify_invocation.UploadFileRequest]}, + {"fetch_app", `{"app_id":"a"}`, echoDispatch[dify_invocation.FetchAppRequest]}, + } + require.Len(t, dispatchMapping, len(cases)) + for _, tc := range cases { + t.Run(string(tc.kind), func(t *testing.T) { + require.Contains(t, dispatchMapping, tc.kind) + replaceDispatch(t, tc.kind, tc.dispatch) + response := invokeJSONRequest(t, tc.kind, tc.payload) + if tc.kind == "tool" || tc.kind == "app" { + require.Contains(t, string(response), `"id":9007199254740993`) + } + }) + } +} + +func TestInvokeDifyRejectsInvalidRequests(t *testing.T) { + replaceDispatch(t, dify_invocation.INVOKE_TYPE_LLM, echoDispatch[dify_invocation.InvokeLLMRequest]) + for _, wire := range []string{ + `{`, `null`, `[]`, `{}`, `{"type":null}`, `{"type":123}`, `{"type":"llm"}`, + `{"type":"llm","backwards_request_id":null,"request":{}}`, + `{"type":"llm","backwards_request_id":"r","request":null}`, + `{"type":"llm","backwards_request_id":"r","request":[]}`, + `{"type":"llm","backwards_request_id":"r","request":{}} {}`, + } { + t.Run(wire, func(t *testing.T) { + _, err := invokeJSON(t, wire, access_types.PLUGIN_ACCESS_TYPE_ENDPOINT, true) + require.Error(t, err) + }) + } + for _, tc := range []struct { + name string + wire string + from access_types.PluginAccessType + allowed bool + message string + }{ + {"required fields", `{"type":"llm","backwards_request_id":"r","request":{}}`, access_types.PLUGIN_ACCESS_TYPE_ENDPOINT, true, "error validating struct"}, + {"invalid content", `{"type":"llm","backwards_request_id":"r","request":{"provider":"p","model":"m","mode":"chat","prompt_messages":[{"role":"assistant","content":123}]}}`, access_types.PLUGIN_ACCESS_TYPE_ENDPOINT, true, "content must be a string or an array"}, + {"permissions", `{"type":"llm","backwards_request_id":"r","request":{}}`, access_types.PLUGIN_ACCESS_TYPE_ENDPOINT, false, "permission denied"}, + {"model origin", `{"type":"llm","backwards_request_id":"r","request":{}}`, access_types.PLUGIN_ACCESS_TYPE_MODEL, true, "you can not invoke dify"}, + {"unsupported type", `{"type":"unsupported","backwards_request_id":"r","request":{}}`, access_types.PLUGIN_ACCESS_TYPE_ENDPOINT, true, "unsupported invoke type"}, + } { + t.Run(tc.name, func(t *testing.T) { + events, err := invokeJSON(t, tc.wire, tc.from, tc.allowed) + require.NoError(t, err) + require.Len(t, events, 2) + require.Equal(t, REQUEST_EVENT_ERROR, events[0].Event) + require.Contains(t, events[0].Message, tc.message) + }) + } +} diff --git a/internal/core/io_tunnel/backwards_invocation/request.go b/internal/core/io_tunnel/backwards_invocation/request.go index 9e51ebe25..29e5aa291 100644 --- a/internal/core/io_tunnel/backwards_invocation/request.go +++ b/internal/core/io_tunnel/backwards_invocation/request.go @@ -2,6 +2,7 @@ package backwards_invocation import ( "context" + "encoding/json" "fmt" "github.com/langgenius/dify-plugin-daemon/pkg/utils/log" @@ -26,7 +27,7 @@ type BackwardsInvocationWriter interface { type BackwardsInvocation struct { typ BackwardsInvocationType id string - detailedRequest map[string]any + detailedRequest map[string]json.RawMessage session *session_manager.Session // writer is the writer that writes the data to the session @@ -42,7 +43,7 @@ func NewBackwardsInvocation( id string, session *session_manager.Session, writer BackwardsInvocationWriter, - detailedRequest map[string]any, + detailedRequest map[string]json.RawMessage, ) *BackwardsInvocation { return &BackwardsInvocation{ typ: typ, @@ -93,7 +94,7 @@ func (bi *BackwardsInvocation) Type() BackwardsInvocationType { return bi.typ } -func (bi *BackwardsInvocation) RequestData() map[string]any { +func (bi *BackwardsInvocation) RequestData() map[string]json.RawMessage { return bi.detailedRequest } diff --git a/internal/core/io_tunnel/backwards_invocation/task.go b/internal/core/io_tunnel/backwards_invocation/task.go index 6d09c4e69..062dea592 100644 --- a/internal/core/io_tunnel/backwards_invocation/task.go +++ b/internal/core/io_tunnel/backwards_invocation/task.go @@ -1,9 +1,12 @@ package backwards_invocation import ( + "bytes" "encoding/hex" + "encoding/json" "errors" "fmt" + "strings" "github.com/langgenius/dify-plugin-daemon/internal/core/dify_invocation" "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel/access_types" @@ -11,8 +14,8 @@ import ( "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" routinepkg "github.com/langgenius/dify-plugin-daemon/pkg/routine" - "github.com/langgenius/dify-plugin-daemon/pkg/utils/parser" "github.com/langgenius/dify-plugin-daemon/pkg/utils/routine" + "github.com/langgenius/dify-plugin-daemon/pkg/validators" ) // returns error only if payload is not correct @@ -23,25 +26,34 @@ func InvokeDify( writer BackwardsInvocationWriter, data []byte, ) error { - // unmarshal invoke data - request, err := parser.UnmarshalJsonBytes2Map(data) - if err != nil { + // Keep payload values as JSON until typed dispatch, including opaque data and integers. + var request struct { + Type *BackwardsInvocationType `json:"type"` + BackwardsRequestID *string `json:"backwards_request_id"` + Request map[string]json.RawMessage `json:"request"` + } + if err := json.Unmarshal(data, &request); err != nil { return fmt.Errorf("unmarshal invoke request failed: %s", err.Error()) } - if request == nil { - return fmt.Errorf("invoke request is empty") + if request.Type == nil { + return errors.New("invoke request missing type") + } + if request.BackwardsRequestID == nil { + return errors.New("invoke request missing request_id") + } + if request.Request == nil { + return errors.New("invoke request missing request") } // prepare invocation arguments - requestHandle, err := prepareDifyInvocationArguments( + requestHandle := NewBackwardsInvocation( + *request.Type, + *request.BackwardsRequestID, session, writer, - request, + request.Request, ) - if err != nil { - return err - } if invoke_from == access_types.PLUGIN_ACCESS_TYPE_MODEL { requestHandle.WriteError(fmt.Errorf("you can not invoke dify from %s", invoke_from)) @@ -61,8 +73,8 @@ func InvokeDify( routinepkg.RoutineLabelKeyModule: "plugin_daemon", routinepkg.RoutineLabelKeyMethod: "InvokeDify", }, func() { - dispatchDifyInvocationTask(requestHandle) defer requestHandle.EndResponse() + dispatchDifyInvocationTask(requestHandle) }) return nil @@ -195,37 +207,6 @@ func checkPermission(runtime *plugin_entities.PluginDeclaration, requestHandle * return nil } -func prepareDifyInvocationArguments( - session *session_manager.Session, - writer BackwardsInvocationWriter, - request map[string]any, -) (*BackwardsInvocation, error) { - typ, ok := request["type"].(string) - if !ok { - return nil, fmt.Errorf("invoke request missing type: %s", request) - } - - // get request id - backwardsRequestId, ok := request["backwards_request_id"].(string) - if !ok { - return nil, fmt.Errorf("invoke request missing request_id: %s", request) - } - - // get request - detailedRequest, ok := request["request"].(map[string]any) - if !ok { - return nil, fmt.Errorf("invoke request missing request: %s", request) - } - - return NewBackwardsInvocation( - BackwardsInvocationType(typ), - backwardsRequestId, - session, - writer, - detailedRequest, - ), nil -} - var ( dispatchMapping = map[dify_invocation.InvokeType]func(handle *BackwardsInvocation){ dify_invocation.INVOKE_TYPE_TOOL: func(handle *BackwardsInvocation) { @@ -289,12 +270,23 @@ func genericDispatchTask[T any]( request *T, ), ) { - r, err := parser.MapToStruct[T](handle.RequestData()) + data, err := json.Marshal(handle.RequestData()) if err != nil { + handle.WriteError(fmt.Errorf("marshal backwards invoke request failed: %s", err.Error())) + return + } + var request T + decoder := json.NewDecoder(bytes.NewReader(data)) + decoder.UseNumber() + if err := decoder.Decode(&request); err != nil { handle.WriteError(fmt.Errorf("unmarshal backwards invoke request failed: %s", err.Error())) return } - dispatch(handle, r) + if err := validators.GlobalEntitiesValidator.Struct(request); err != nil { + handle.WriteError(fmt.Errorf("unmarshal backwards invoke request failed: error validating struct: %s", err.Error())) + return + } + dispatch(handle, &request) } func dispatchDifyInvocationTask(handle *BackwardsInvocation) { @@ -304,15 +296,21 @@ func dispatchDifyInvocationTask(handle *BackwardsInvocation) { handle.WriteError(fmt.Errorf("get tenant id failed: %s", err.Error())) return } - requestData["tenant_id"] = tenantId userId, err := handle.UserID() if err != nil { handle.WriteError(fmt.Errorf("get user id failed: %s", err.Error())) return } - requestData["user_id"] = userId + // JSON also matches case-folded field names; remove aliases before injecting trusted identity. + for key := range requestData { + if strings.EqualFold(key, "tenant_id") || strings.EqualFold(key, "user_id") || strings.EqualFold(key, "type") { + delete(requestData, key) + } + } + requestData["tenant_id"], _ = json.Marshal(tenantId) + requestData["user_id"], _ = json.Marshal(userId) typ := handle.Type() - requestData["type"] = typ + requestData["type"], _ = json.Marshal(typ) for t, v := range dispatchMapping { if t == handle.Type() { diff --git a/internal/core/io_tunnel/backwards_invocation/transaction/json_transport_test.go b/internal/core/io_tunnel/backwards_invocation/transaction/json_transport_test.go new file mode 100644 index 000000000..cdb0142f4 --- /dev/null +++ b/internal/core/io_tunnel/backwards_invocation/transaction/json_transport_test.go @@ -0,0 +1,96 @@ +package transaction + +import ( + "bytes" + "context" + "encoding/json" + "testing" + "time" + + "github.com/langgenius/dify-plugin-daemon/internal/core/dify_invocation" + "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel/access_types" + "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel/backwards_invocation" + "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" + "github.com/langgenius/dify-plugin-daemon/pkg/entities/model_entities" + "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/parser" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/routine" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" + "github.com/stretchr/testify/require" +) + +type replayInvocation struct { + dify_invocation.BackwardsInvocation + request *dify_invocation.InvokeLLMRequest +} + +func (i *replayInvocation) SetContext(context.Context) {} + +func (i *replayInvocation) InvokeLLM(request *dify_invocation.InvokeLLMRequest) (*stream.Stream[model_entities.LLMResultChunk], error) { + i.request = request + result := stream.NewStream[model_entities.LLMResultChunk](1) + index := 0 + result.Write(model_entities.LLMResultChunk{ + Model: "model", + Delta: model_entities.LLMResultChunkDelta{Index: &index, Message: request.PromptMessages[0]}, + }) + result.Close() + return result, nil +} + +type replayResponseWriter struct { + bytes.Buffer + done chan struct{} +} + +func (w *replayResponseWriter) Flush() {} +func (w *replayResponseWriter) Close() error { close(w.done); return nil } + +func TestServerlessInvocationPreservesRawJSON(t *testing.T) { + routine.InitPool(4) + invocation := &replayInvocation{} + session := session_manager.NewSession(session_manager.NewSessionPayload{ + TenantID: "tenant", UserID: "user", BackwardsInvocation: invocation, IgnoreCache: true, + }) + t.Cleanup(func() { + session_manager.DeleteSession(session_manager.DeleteSessionPayload{ID: session.ID, IgnoreCache: true}) + }) + declaration := &plugin_entities.PluginDeclaration{ + PluginDeclarationWithoutAdvancedFields: plugin_entities.PluginDeclarationWithoutAdvancedFields{ + Resource: plugin_entities.PluginResourceRequirement{ + Permission: &plugin_entities.PluginPermissionRequirement{ + Model: &plugin_entities.PluginPermissionModelRequirement{Enabled: true, LLM: true}, + }, + }, + }, + } + writer := &replayResponseWriter{done: make(chan struct{})} + body, err := json.Marshal(plugin_entities.PluginUniversalEvent{ + SessionId: session.ID, + Event: plugin_entities.PLUGIN_EVENT_SESSION, + Data: json.RawMessage(`{"type":"invoke","data":{"type":"llm","backwards_request_id":"request","request":{"provider":"p","model":"m","mode":"chat","completion_params":{"seed":9007199254740993},"prompt_messages":[{"role":"assistant","content":null,"opaque_body":{"id":9007199254740993}}]}}}`), + }) + require.NoError(t, err) + // Follow the serverless handler's event parsing and invoke path with a controlled session. + plugin_entities.ParsePluginUniversalEvent(body, "", func(sessionID string, data []byte) { + require.Equal(t, session.ID, sessionID) + message, err := parser.UnmarshalJsonBytes[plugin_entities.SessionMessage](data) + require.NoError(t, err) + require.NoError(t, backwards_invocation.InvokeDify( + declaration, access_types.PLUGIN_ACCESS_TYPE_ENDPOINT, session, + NewServerlessTransactionWriter(session, writer), message.Data, + )) + }, nil, func(err string) { + t.Fatalf("unexpected event error: %s", err) + }, nil) + select { + case <-writer.done: + case <-time.After(5 * time.Second): + t.Fatal("serverless writer did not close") + } + require.NotNil(t, invocation.request) + require.Equal(t, json.Number("9007199254740993"), invocation.request.CompletionParams["seed"]) + require.Equal(t, `{"id":9007199254740993}`, string(invocation.request.PromptMessages[0].OpaqueBody)) + require.Contains(t, writer.String(), `"opaque_body":{"id":9007199254740993}`) + require.Contains(t, writer.String(), `"event":"end"`) +} diff --git a/internal/core/io_tunnel/backwards_invocation/transaction/serverless_handler.go b/internal/core/io_tunnel/backwards_invocation/transaction/serverless_handler.go index 2c5414131..f67f62ce0 100644 --- a/internal/core/io_tunnel/backwards_invocation/transaction/serverless_handler.go +++ b/internal/core/io_tunnel/backwards_invocation/transaction/serverless_handler.go @@ -1,6 +1,7 @@ package transaction import ( + "encoding/json" "fmt" "io" "net/http" @@ -148,10 +149,13 @@ func (h *ServerlessTransactionHandler) Handle(ctx *gin.Context, sessionId string session, err := session_manager.GetSession(sessionId) if err != nil { ctx.Writer.WriteHeader(http.StatusBadRequest) - invokePayload, marshalErr := parser.UnmarshalJsonBytes2Map(sessionMessage.Data) + var invokePayload struct { + BackwardsRequestID string `json:"backwards_request_id"` + } + marshalErr := json.Unmarshal(sessionMessage.Data, &invokePayload) var backwardsRequestId string if marshalErr == nil { - backwardsRequestId, _ = invokePayload["backwards_request_id"].(string) + backwardsRequestId = invokePayload.BackwardsRequestID } cacheErrorKind := session_manager.SessionCacheErrorKind(err) diff --git a/internal/core/io_tunnel/backwards_invocation/transaction/serverless_handler_test.go b/internal/core/io_tunnel/backwards_invocation/transaction/serverless_handler_test.go index da7f1c373..903210456 100644 --- a/internal/core/io_tunnel/backwards_invocation/transaction/serverless_handler_test.go +++ b/internal/core/io_tunnel/backwards_invocation/transaction/serverless_handler_test.go @@ -28,6 +28,7 @@ func TestHandle_SessionNotFound_WritesErrorResponse(t *testing.T) { invokePayload := map[string]any{ "backwards_request_id": testBackwardsReqID, + "request": json.RawMessage(`{"opaque_body":{"id":1e400}}`), } invokePayloadBytes := parser.MarshalJsonBytes(invokePayload) diff --git a/pkg/entities/model_entities/llm.go b/pkg/entities/model_entities/llm.go index c297d9771..158389ff9 100644 --- a/pkg/entities/model_entities/llm.go +++ b/pkg/entities/model_entities/llm.go @@ -54,7 +54,7 @@ func isPromptMessageRole(fl validator.FieldLevel) bool { type PromptMessage struct { Role PromptMessageRole `json:"role" validate:"required,prompt_message_role"` - Content any `json:"content" validate:"required,prompt_message_content"` + Content any `json:"content" validate:"omitempty,prompt_message_content"` Name string `json:"name"` ToolCalls []PromptMessageToolCall `json:"tool_calls" validate:"dive"` ToolCallId string `json:"tool_call_id"` @@ -145,7 +145,7 @@ func validateToolChoice(sl validator.StructLevel) { func unmarshalPromptMessageContent(data json.RawMessage) (any, error) { trimmed := bytes.TrimSpace(data) if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) { - return nil, errors.New("content field is required") + return nil, nil } switch trimmed[0] { @@ -184,10 +184,6 @@ func (p *PromptMessage) UnmarshalJSON(data []byte) error { if raw.Role == "" { return errors.New("role field is required") } - if len(raw.Content) == 0 { - return errors.New("content field is required") - } - content, err := unmarshalPromptMessageContent(raw.Content) if err != nil { return err diff --git a/pkg/entities/model_entities/llm_test.go b/pkg/entities/model_entities/llm_test.go index 341dbe8d6..1722df50b 100644 --- a/pkg/entities/model_entities/llm_test.go +++ b/pkg/entities/model_entities/llm_test.go @@ -496,3 +496,24 @@ func TestPromptMessageContentWithoutFilename(t *testing.T) { t.Errorf("expected empty filename, got '%s'", content[0].Filename) } } + +func TestOptionalPromptMessageContent(t *testing.T) { + for _, content := range []string{"", `,"content":null`, `,"content":""`, `,"content":[]`} { + t.Run(content, func(t *testing.T) { + wire := `{"role":"assistant","opaque_body":{"signature":"original"},"tool_calls":[{"id":"call","type":"function","function":{"name":"f","arguments":"{}"}}]` + content + `}` + message, err := parser.UnmarshalJsonBytes[PromptMessage]([]byte(wire)) + assert.NoError(t, err) + assert.Equal(t, `{"signature":"original"}`, string(message.OpaqueBody)) + assert.Len(t, message.ToolCalls, 1) + chunk := `{"model":"m","delta":{"index":0,"message":` + wire + `}}` + _, err = parser.UnmarshalJsonBytes[LLMResultChunk]([]byte(chunk)) + assert.NoError(t, err) + }) + } + for _, content := range []string{`123`, `false`, `{}`, `[123]`, `[{"type":"invalid","data":"text"}]`} { + t.Run(content, func(t *testing.T) { + _, err := parser.UnmarshalJsonBytes[PromptMessage]([]byte(`{"role":"assistant","content":` + content + `}`)) + assert.Error(t, err) + }) + } +} From c0f1284266addc70940e025d7f35aecf2f306994 Mon Sep 17 00:00:00 2001 From: WH-2099 Date: Wed, 21 Oct 2026 10:07:46 +0000 Subject: [PATCH 2/2] fix(invocation): preserve numbers in agent requests and responses --- .../io_tunnel/agent_json_transport_test.go | 86 +++++++++++++++++++ internal/core/io_tunnel/agent_service.go | 10 ++- .../server/controllers/json_transport_test.go | 38 ++++++++ pkg/entities/plugin_entities/request.go | 13 +++ pkg/entities/tool_entities/tool.go | 10 +++ 5 files changed, 154 insertions(+), 3 deletions(-) create mode 100644 internal/core/io_tunnel/agent_json_transport_test.go create mode 100644 internal/server/controllers/json_transport_test.go diff --git a/internal/core/io_tunnel/agent_json_transport_test.go b/internal/core/io_tunnel/agent_json_transport_test.go new file mode 100644 index 000000000..adbb11597 --- /dev/null +++ b/internal/core/io_tunnel/agent_json_transport_test.go @@ -0,0 +1,86 @@ +package io_tunnel + +import ( + "encoding/json" + "fmt" + "testing" + + "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel/access_types" + "github.com/langgenius/dify-plugin-daemon/pkg/entities" + "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" + "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" + "github.com/langgenius/dify-plugin-daemon/pkg/entities/tool_entities" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/routine" + "github.com/stretchr/testify/require" +) + +func TestAgentPreservesHistoryAndBlobResponses(t *testing.T) { + routine.InitPool(4) + for _, runtimeType := range []plugin_entities.PluginRuntimeType{ + plugin_entities.PLUGIN_RUNTIME_TYPE_LOCAL, + plugin_entities.PLUGIN_RUNTIME_TYPE_REMOTE, + plugin_entities.PLUGIN_RUNTIME_TYPE_SERVERLESS, + } { + t.Run(string(runtimeType), func(t *testing.T) { + listener := entities.NewCallbackHandler[plugin_entities.SessionMessage]() + runtime := &fakePluginRuntime{ + listener: listener, uniqueIdentifier: testPluginUniqueIdentifier, runtimeType: runtimeType, + } + runtime.writeFn = func(_ string, _ access_types.PluginAccessAction, _ []byte) error { + for _, payload := range []string{ + `{"type":"json","message":{"json_object":{"model_history":{"n":9007199254740993,"state":[false,0,{},[]]}}}}`, + `{"type":"blob_chunk","message":{"id":"file","total_length":2,"blob":"aGk=","end":false}}`, + `{"type":"blob_chunk","message":{"id":"file","total_length":2,"blob":"","end":true}}`, + } { + listener.Send(plugin_entities.SessionMessage{Type: plugin_entities.SESSION_MESSAGE_TYPE_STREAM, Data: []byte(payload)}) + } + listener.Send(plugin_entities.SessionMessage{Type: plugin_entities.SESSION_MESSAGE_TYPE_END}) + return nil + } + session := newTestSession(t, runtime, access_types.PLUGIN_ACCESS_TYPE_AGENT_STRATEGY, + access_types.PLUGIN_ACCESS_ACTION_INVOKE_AGENT_STRATEGY) + response, err := InvokeAgentStrategy(session, &requests.RequestInvokeAgentStrategy{}) + require.NoError(t, err) + defer response.Close() + require.True(t, response.Next()) + history, err := response.Read() + require.NoError(t, err) + encoded, err := json.Marshal(history) + require.NoError(t, err) + require.Contains(t, string(encoded), `"n":9007199254740993`) + require.Contains(t, string(encoded), `"state":[false,0,{},[]]`) + require.True(t, response.Next()) + blob, err := response.Read() + require.NoError(t, err) + require.Equal(t, tool_entities.ToolResponseChunkTypeBlob, blob.Type) + require.Equal(t, []byte("hi"), blob.Message["blob"]) + require.False(t, response.Next()) + }) + } +} + +func TestAgentRejectsInvalidBlobLengths(t *testing.T) { + routine.InitPool(4) + for _, length := range []string{"-1", "15728641", "1.5", "9223372036854775808"} { + t.Run(length, func(t *testing.T) { + listener := entities.NewCallbackHandler[plugin_entities.SessionMessage]() + runtime := &fakePluginRuntime{listener: listener, uniqueIdentifier: testPluginUniqueIdentifier} + runtime.writeFn = func(_ string, _ access_types.PluginAccessAction, _ []byte) error { + listener.Send(plugin_entities.SessionMessage{ + Type: plugin_entities.SESSION_MESSAGE_TYPE_STREAM, + Data: []byte(fmt.Sprintf(`{"type":"blob_chunk","message":{"id":"file","total_length":%s,"blob":"","end":false}}`, length)), + }) + listener.Send(plugin_entities.SessionMessage{Type: plugin_entities.SESSION_MESSAGE_TYPE_END}) + return nil + } + session := newTestSession(t, runtime, access_types.PLUGIN_ACCESS_TYPE_AGENT_STRATEGY, + access_types.PLUGIN_ACCESS_ACTION_INVOKE_AGENT_STRATEGY) + response, err := InvokeAgentStrategy(session, &requests.RequestInvokeAgentStrategy{}) + require.NoError(t, err) + defer response.Close() + require.True(t, response.Next()) + _, err = response.Read() + require.EqualError(t, err, "invalid file length") + }) + } +} diff --git a/internal/core/io_tunnel/agent_service.go b/internal/core/io_tunnel/agent_service.go index 735d97604..d87c8a27b 100644 --- a/internal/core/io_tunnel/agent_service.go +++ b/internal/core/io_tunnel/agent_service.go @@ -3,6 +3,7 @@ package io_tunnel import ( "bytes" "encoding/base64" + "encoding/json" "errors" "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" @@ -57,13 +58,16 @@ func InvokeAgentStrategy( continue } - totalLength, ok := item.Message["total_length"].(float64) + totalLength, ok := item.Message["total_length"].(json.Number) if !ok { continue } - // convert total_length to int - totalLengthInt := int(totalLength) + totalLengthInt, err := totalLength.Int64() + if err != nil || totalLengthInt < 0 || totalLengthInt > 15*1024*1024 { + newResponse.WriteError(errors.New("invalid file length")) + return + } blob, ok := item.Message["blob"].(string) if !ok { diff --git a/internal/server/controllers/json_transport_test.go b/internal/server/controllers/json_transport_test.go new file mode 100644 index 000000000..98d98b784 --- /dev/null +++ b/internal/server/controllers/json_transport_test.go @@ -0,0 +1,38 @@ +package controllers + +import ( + "encoding/json" + "net/http/httptest" + "strings" + "testing" + + "github.com/gin-gonic/gin" + "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" + "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/parser" + "github.com/stretchr/testify/require" +) + +func checkDispatchJSON[T any](t *testing.T, payload string) { + t.Helper() + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest("POST", "/", strings.NewReader(payload)) + c.Request.Header.Set("Content-Type", "application/json") + called := false + BindRequest(c, func(request plugin_entities.InvokePluginRequest[T]) { + called = true + encoded, err := json.Marshal(parser.StructToMap(request.Data)) + require.NoError(t, err) + require.Contains(t, string(encoded), `"n":9007199254740993`) + require.Contains(t, string(encoded), `"state":[false,0,{},[]]`) + }) + require.True(t, called) +} + +func TestDispatchPreservesToolArguments(t *testing.T) { + checkDispatchJSON[requests.RequestInvokeTool](t, `{"tenant_id":"t","data":{"provider":"p","tool":"lookup","tool_parameters":{"n":9007199254740993,"state":[false,0,{},[]]}}}`) +} + +func TestDispatchPreservesAgentHistory(t *testing.T) { + checkDispatchJSON[requests.RequestInvokeAgentStrategy](t, `{"tenant_id":"t","data":{"agent_strategy_provider":"p","agent_strategy":"fc","agent_strategy_params":{"model":{"history_prompt_messages":[{"role":"assistant","opaque_body":{"n":9007199254740993,"state":[false,0,{},[]]}}]}}}}`) +} diff --git a/pkg/entities/plugin_entities/request.go b/pkg/entities/plugin_entities/request.go index 969ead608..f0bc865dd 100644 --- a/pkg/entities/plugin_entities/request.go +++ b/pkg/entities/plugin_entities/request.go @@ -1,5 +1,10 @@ package plugin_entities +import ( + "bytes" + "encoding/json" +) + type InvokePluginUserIdentity struct { TenantId string `json:"tenant_id" validate:"required" uri:"tenant_id"` UserId string `json:"user_id"` @@ -22,3 +27,11 @@ type InvokePluginRequest[T any] struct { Data T `json:"data" validate:"required"` } + +func (r *InvokePluginRequest[T]) UnmarshalJSON(data []byte) error { + type request InvokePluginRequest[T] + decoder := json.NewDecoder(bytes.NewReader(data)) + // Plugin arguments and saved model state must retain their original JSON numbers. + decoder.UseNumber() + return decoder.Decode((*request)(r)) +} diff --git a/pkg/entities/tool_entities/tool.go b/pkg/entities/tool_entities/tool.go index 784ad4641..a70479866 100644 --- a/pkg/entities/tool_entities/tool.go +++ b/pkg/entities/tool_entities/tool.go @@ -1,6 +1,9 @@ package tool_entities import ( + "bytes" + "encoding/json" + "github.com/go-playground/validator/v10" "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" "github.com/langgenius/dify-plugin-daemon/pkg/validators" @@ -57,6 +60,13 @@ type ToolResponseChunk struct { Meta map[string]any `json:"meta"` } +func (r *ToolResponseChunk) UnmarshalJSON(data []byte) error { + type response ToolResponseChunk + decoder := json.NewDecoder(bytes.NewReader(data)) + decoder.UseNumber() + return decoder.Decode((*response)(r)) +} + type GetToolRuntimeParametersResponse struct { Parameters []plugin_entities.ToolParameter `json:"parameters"` }