ragflow/internal/agent/sandbox/result_protocol_test.go

//
//  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
//

package sandbox

import (
	"encoding/base64"
	"encoding/json"
	"strings"
	"testing"
)

func TestBuildPythonWrapper_ContainsMainAndArgs(t *testing.T) {
	t.Parallel()
	wrapped := BuildPythonWrapper("def main(x): return x * 2", `{"x": 21}`)
	if !strings.Contains(wrapped, "def main(x): return x * 2") {
		t.Errorf("wrapper missing user code; got: %s", wrapped)
	}
	// argsJSON is base64-encoded in the wrapper, so the call site
	// must use json.loads(base64.b64decode(...)).
	if !strings.Contains(wrapped, "main(**json.loads(base64.b64decode(") {
		t.Errorf("wrapper missing main(**json.loads(base64.b64decode(...))) call; got: %s", wrapped)
	}
	if !strings.Contains(wrapped, resultMarkerPrefix) {
		t.Errorf("wrapper missing result marker; got: %s", wrapped)
	}
	// The marker line carries a base64-encoded JSON payload. We do
	// not decode it here — the round-trip is exercised in
	// TestExtractStructuredResult_RoundTrip.
}

func TestBuildJavaScriptWrapper_ContainsMainAndArgs(t *testing.T) {
	t.Parallel()
	wrapped := BuildJavaScriptWrapper("async function main(args) { return args.x * 2 }", `{"x": 21}`)
	if !strings.Contains(wrapped, "async function main(args)") {
		t.Errorf("wrapper missing user code; got: %s", wrapped)
	}
	if !strings.Contains(wrapped, "const __ragflowArgsB64 = ") {
		t.Errorf("wrapper missing base64 args literal; got: %s", wrapped)
	}
	if !strings.Contains(wrapped, "JSON.parse(Buffer.from(__ragflowArgsB64") {
		t.Errorf("wrapper missing args decoding; got: %s", wrapped)
	}
	if !strings.Contains(wrapped, resultMarkerPrefix) {
		t.Errorf("wrapper missing result marker; got: %s", wrapped)
	}
}

func TestExtractStructuredResult_RoundTrip(t *testing.T) {
	t.Parallel()

	// Simulate what the wrapper prints: stdout with a marker
	// carrying a base64 JSON payload, then any trailing output.
	payload := map[string]any{"present": true, "value": 42, "type": "json"}
	raw, _ := json.Marshal(payload)
	markerLine := resultMarkerPrefix + base64.StdEncoding.EncodeToString(raw)
	stdout := "Hello, world!\n" + markerLine + "\nTrailing line\n"

	cleaned, structured := ExtractStructuredResult(stdout)
	if structured == nil {
		t.Fatalf("structured result is nil; cleaned=%q", cleaned)
	}
	if v, ok := structured["value"].(float64); !ok || v != 42 {
		t.Errorf("structured[value] = %v, want 42", structured["value"])
	}
	if strings.Contains(cleaned, resultMarkerPrefix) {
		t.Errorf("cleaned stdout still contains marker: %q", cleaned)
	}
	if !strings.Contains(cleaned, "Hello, world!") || !strings.Contains(cleaned, "Trailing line") {
		t.Errorf("cleaned stdout lost user lines: %q", cleaned)
	}
}

func TestExtractStructuredResult_NoMarker(t *testing.T) {
	t.Parallel()
	stdout := "Plain output\nNothing fancy"
	cleaned, structured := ExtractStructuredResult(stdout)
	if cleaned != stdout {
		t.Errorf("cleaned = %q, want %q", cleaned, stdout)
	}
	if len(structured) != 0 {
		t.Errorf("structured = %v, want empty map", structured)
	}
}

func TestExtractStructuredResult_Empty(t *testing.T) {
	t.Parallel()
	cleaned, structured := ExtractStructuredResult("")
	if cleaned != "" {
		t.Errorf("cleaned = %q, want empty", cleaned)
	}
	if structured == nil {
		t.Errorf("structured is nil, want empty map (not nil)")
	}
	if len(structured) != 0 {
		t.Errorf("structured = %v, want empty", structured)
	}
}

func TestExtractStructuredResult_UndecodableMarker(t *testing.T) {
	t.Parallel()
	stdout := "good line\n" + resultMarkerPrefix + "!!!not-base64!!!\nmore output"
	cleaned, structured := ExtractStructuredResult(stdout)
	// Undecodable marker lines are kept in cleaned stdout (Python
	// behavior: except branch appends the line back). The user can
	// see the raw base64 and debug.
	if !strings.Contains(cleaned, "!!!not-base64!!!") {
		t.Errorf("cleaned dropped the undecodable marker line: %q", cleaned)
	}
	if len(structured) != 0 {
		t.Errorf("structured = %v, want empty", structured)
	}
}

func TestExtractStructuredResult_MultipleMarkers_LastWins(t *testing.T) {
	t.Parallel()
	// Python implementation overwrites structured_result on each
	// marker line, so the last marker wins.
	payload1 := map[string]any{"present": true, "value": "first", "type": "json"}
	payload2 := map[string]any{"present": true, "value": "second", "type": "json"}
	raw1, _ := json.Marshal(payload1)
	raw2, _ := json.Marshal(payload2)
	stdout := resultMarkerPrefix + base64.StdEncoding.EncodeToString(raw1) + "\n" +
		resultMarkerPrefix + base64.StdEncoding.EncodeToString(raw2) + "\n"
	_, structured := ExtractStructuredResult(stdout)
	if v, _ := structured["value"].(string); v != "second" {
		t.Errorf("structured[value] = %v, want 'second'", structured["value"])
	}
}

func TestExtractStructuredResult_PreservesTrailingNewline(t *testing.T) {
	t.Parallel()
	payload := map[string]any{"present": true, "value": 1, "type": "json"}
	raw, _ := json.Marshal(payload)
	stdout := "line\n" + resultMarkerPrefix + base64.StdEncoding.EncodeToString(raw) + "\n"
	cleaned, _ := ExtractStructuredResult(stdout)
	if !strings.HasSuffix(cleaned, "\n") {
		t.Errorf("cleaned lost trailing newline: %q", cleaned)
	}
}

func TestArgsToJSON(t *testing.T) {
	t.Parallel()

	t.Run("nil map", func(t *testing.T) {
		s, err := argsToJSON(nil)
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		if s != "{}" {
			t.Errorf("got %q, want %q", s, "{}")
		}
	})

	t.Run("empty map", func(t *testing.T) {
		s, err := argsToJSON(map[string]any{})
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		if s != "{}" {
			t.Errorf("got %q, want %q", s, "{}")
		}
	})

	t.Run("populated map", func(t *testing.T) {
		s, err := argsToJSON(map[string]any{"x": 1, "y": "hi"})
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		// The exact JSON key order is non-deterministic; check both
		// fields independently.
		var got map[string]any
		if err := json.Unmarshal([]byte(s), &got); err != nil {
			t.Fatalf("invalid JSON: %v", err)
		}
		if got["x"].(float64) != 1 {
			t.Errorf("x = %v, want 1", got["x"])
		}
		if got["y"].(string) != "hi" {
			t.Errorf("y = %v, want 'hi'", got["y"])
		}
	})
}

func TestNormalizeLanguage(t *testing.T) {
	t.Parallel()

	cases := []struct {
		in, want string
	}{
		{"python", "python"},
		{"python3", "python"},
		{"Python", "python"},
		{"javascript", "nodejs"},
		{"js", "nodejs"},
		{"nodejs", "nodejs"},
		{"node", "nodejs"},
		{"", ""},
		{"ruby", ""},
	}
	for _, tc := range cases {
		if got := normalizeLanguage(tc.in); got != tc.want {
			t.Errorf("normalizeLanguage(%q) = %q, want %q", tc.in, got, tc.want)
		}
	}
}

func TestValidateTimeout(t *testing.T) {
	t.Parallel()

	if _, err := validateTimeout(0); err == nil {
		t.Errorf("validateTimeout(0) returned no error, want one")
	}
	if _, err := validateTimeout(-1); err == nil {
		t.Errorf("validateTimeout(-1) returned no error, want one")
	}
	if _, err := validateTimeout(700); err == nil {
		t.Errorf("validateTimeout(700) returned no error, want one (max is 600)")
	}
	if got, err := validateTimeout(30); err != nil || got != 30 {
		t.Errorf("validateTimeout(30) = %d, %v; want 30, nil", got, err)
	}
	if got, err := validateTimeout(600); err != nil || got != 600 {
		t.Errorf("validateTimeout(600) = %d, %v; want 600, nil", got, err)
	}
}