Add files via upload

internal/security/executor_test.go

@@ -94,6 +94,60 @@ func TestExecuteSystemCommand_FailureFormat(t *testing.T) {
 	}
 }
 
+func TestExecuteSystemCommand_OutputIsSourceLimited(t *testing.T) {
+	executor, _ := setupTestExecutor(t)
+	executor.SetToolOutputMaxBytes(64)
+	res, err := executor.executeSystemCommand(context.Background(), map[string]interface{}{
+		"command": "i=0; while [ $i -lt 2000 ]; do printf 0123456789; i=$((i+1)); done",
+		"shell":   "sh",
+	})
+	if err != nil {
+		t.Fatalf("executeSystemCommand: %v", err)
+	}
+	if res == nil || res.IsError {
+		t.Fatalf("expected success, got %+v", res)
+	}
+	text := res.Content[0].Text
+	if !strings.Contains(text, "tool output limit reached") {
+		t.Fatalf("missing output limit marker: %q", text)
+	}
+	if len(text) > 64 {
+		t.Fatalf("output exceeded hard limit: len=%d text=%q", len(text), text)
+	}
+	if strings.Contains(text, strings.Repeat("0123456789", 20)) {
+		t.Fatalf("output kept too much data: len=%d", len(text))
+	}
+}
+
+func TestExecuteSystemCommand_StreamingOutputIsSourceLimited(t *testing.T) {
+	executor, _ := setupTestExecutor(t)
+	executor.SetToolOutputMaxBytes(64)
+	var streamed strings.Builder
+	ctx := context.WithValue(context.Background(), ToolOutputCallbackCtxKey, ToolOutputCallback(func(chunk string) {
+		streamed.WriteString(chunk)
+	}))
+	res, err := executor.executeSystemCommand(ctx, map[string]interface{}{
+		"command": "i=0; while [ $i -lt 2000 ]; do printf abcdefghij; i=$((i+1)); done",
+		"shell":   "sh",
+	})
+	if err != nil {
+		t.Fatalf("executeSystemCommand: %v", err)
+	}
+	text := res.Content[0].Text
+	if text != streamed.String() {
+		t.Fatalf("streamed output and returned output diverged\nstream=%q\nret=%q", streamed.String(), text)
+	}
+	if !strings.Contains(text, "tool output limit reached") {
+		t.Fatalf("missing output limit marker: %q", text)
+	}
+	if len(text) > 64 {
+		t.Fatalf("streamed output exceeded hard limit: len=%d text=%q", len(text), text)
+	}
+	if strings.Contains(text, strings.Repeat("abcdefghij", 20)) {
+		t.Fatalf("streamed output kept too much data: len=%d", len(text))
+	}
+}
+
 func TestBuildCommandArgs_NmapSkipsEmptyOptionalFlags(t *testing.T) {
 	pos1 := 1
 	executor, _ := setupTestExecutor(t)